#!/usr/bin/env python """ Benchmark and validate SDNQ attention and weight dequantization on the local GPU. The attention section runs the kernel from sdnq/kernels/triton_atten.py directly and compares speed and numerical error against torch scaled_dot_product_attention and sageattention when installed. Verifies mask, causal, GQA, cross-attention and padding code paths, probes float8 hardware support and the torch.compile input prep, and prints recommended values for the Compute Settings -> SDNQ Attention section. The dequant section builds real SDNQ linear layers per storage dtype (int8, uint4, float8_e4m3fn, float8_e4m3fn_sdnq, float4_e2m1fn) at Flux.1, Krea 2 and Qwen3 TE layer geometry and measures eager vs compiled weight dequantization plus the full linear forward with and without quantized matmul, against a bf16 nn.Linear baseline. For float storage dtypes it also measures quantized matmul with the MatMul type set explicitly (int8, float16), since enabled routes them to fp8 matmul, which not every gpu can run. Setting sweeps (--dequant-sweeps) cover group size, svd rank, weight-side hadamard group size, dequantize full precision off, dynamic quantization, cpu quantize time and conv2d quantization. Probes whether compiled dequant of fp8 storage works on this GPU (triton before sm_89 lacks e4m3 conversions) and prints recommended values for the Compute Settings -> SDNQ section. The block section measures complete configurations (weights dtype x matmul path x attention) end to end through a dit-style transformer block, with output error at depth one and four against an fp32 reference block, because component speedups and errors do not compose multiplicatively. Benchmarks run in the webui's configured dtype (--dtype overrides). The prep column is the q/k/v quantization cost outside the kernel, included in the median. Recommendations come from measured comparisons only. Shape presets follow real model geometries: sd15, sdxl, sdxl-cross, qwen3-te (Anima TE), anima, flux2 (Klein), krea2 (segment mask), wan22 (A14B), ltx2 (LTX 2.3), plus masked and wan22-cfg presets. Run from the sdnext root with the venv active: python cli/sdnq-attention-benchmark.py python cli/sdnq-attention-benchmark.py --shapes all --json results.json python cli/sdnq-attention-benchmark.py --sections dequant python cli/sdnq-attention-benchmark.py --shapes wan22,ltx2 --iters 20 The first run of each shape includes triton autotune and torch.compile time; results are cached on disk and reused by the webui for matching shapes. """ import io import os import sys import json import math import time import signal import inspect import logging import argparse import tempfile import statistics import importlib import importlib.metadata from contextlib import contextmanager import torch from rich import box from rich.console import Console, Group from rich.live import Live from rich.markup import escape from rich.panel import Panel from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn from rich.table import Table console = Console(width=None if sys.stdout.isatty() else 140) # keep tables readable when piped to a file ROUNDED_BOX = box.ROUNDED # modules.logger monkeypatches rich's box.ROUNDED to box.SIMPLE at import; keep the real one for panels shared = None devices = None sdnq_triton_atten = None bench_dtype = torch.bfloat16 # resolved to the webui's configured dtype in main transcript = [] # final tables and panels, written out by --save report = {} # structured results mirroring the tables, written out by --json compiled_dequantize_weight = None # tool-owned compiled variant, built on first use torch_device_module = torch.xpu if torch.xpu.is_available() else torch.cuda torch_device = "xpu" if torch.xpu.is_available() else "cuda" def dtype_label(): return str(bench_dtype).replace("torch.", "") def emit(renderable): console.print(renderable) transcript.append(renderable) def save_transcript(path): with open(path, "w", encoding="utf-8") as fh: file_console = Console(file=fh, width=140) # file consoles are non-terminal: markup renders to plain text for renderable in transcript: file_console.print(renderable) console.print(f"results saved to {path}") def key_padding_mask(cfg, device, keep=0.75): # boolean key-padding mask over the kv axis, first keep fraction of keys valid kv_tokens = cfg.get("kv_tokens", cfg["tokens"]) attn_mask = torch.zeros(cfg["batch"], 1, 1, kv_tokens, device=device, dtype=torch.bool) attn_mask[..., :int(kv_tokens * keep)] = True return attn_mask def krea2_segment_mask(cfg, device): # the transformer's segment_mask: text is padded to a fixed 512 tokens ahead of the # image tokens and the padded tail is masked for queries and keys both, so padding # query rows are fully masked and yield nan under sdpa (the model nan_to_num's them) valid = torch.ones(cfg["batch"], cfg["tokens"], device=device, dtype=torch.bool) valid[:, 128:512] = False return valid.unsqueeze(1).unsqueeze(2) * valid.unsqueeze(1).unsqueeze(3) def build_preset_mask(cfg, device): # dense masks come from the preset's mask_fn; the element guard keeps h3-scale presets # from materializing multi-gigabyte masks, those shapes belong to block-granular masks mask_fn = cfg.get("mask_fn") if mask_fn is None: return None attn_mask = mask_fn(cfg, device) if attn_mask is not None and attn_mask.numel() > 2**31: raise ValueError(f"preset dense mask holds {attn_mask.numel():,} elements; this shape needs a block-granular mask, not a token mask") return attn_mask shape_presets = { # geometry from the model transformer and text-encoder configs; # optional keys: kv_tokens (cross-attention), kv_heads (gqa), causal, # mask_fn (token-granular attn_mask builder), mask_nan_guard (fully-masked query # rows nan under stock sdpa), iters/warmup (per-preset run overrides for very large # shapes), ref_head_chunk (head-sliced fp32 reference to bound peak memory), # sparse (token layout driving the sparse selector rows) "sd15": dict(batch=2, heads=8, tokens=4096, head_dim=40, desc="SD 1.5 unet self-attention at 512px, batched cfg, head dim padded 40 to 64"), "sdxl": dict(batch=2, heads=10, tokens=4096, head_dim=64, desc="SDXL unet self-attention at 1024px, batched cfg"), "sdxl-cross": dict(batch=2, heads=10, tokens=4096, kv_tokens=77, head_dim=64, desc="SDXL unet cross-attention at 1024px, 77 text tokens"), "qwen3-te": dict(batch=2, heads=16, kv_heads=8, tokens=512, head_dim=128, causal=True, desc="Qwen3 text encoder (Anima), causal gqa 16:8 heads, 512 token prompt"), "anima": dict(batch=1, heads=16, tokens=4096, head_dim=128, desc="Anima 1.0 self-attention at 1024px, one cfg pass"), "flux2": dict(batch=1, heads=32, tokens=4608, head_dim=128, desc="FLUX.2 Klein 9B joint attention at 1024px, 4096 image plus 512 text tokens"), "krea2": dict(batch=1, heads=48, tokens=4608, head_dim=128, mask_fn=krea2_segment_mask, mask_nan_guard=True, desc="Krea 2 12B joint attention at 1024px, 4096 image plus 512 text tokens (128 real), kv expanded from gqa 48:12, segment mask"), "wan22": dict(batch=1, heads=40, tokens=32760, head_dim=128, desc="Wan 2.2 A14B self-attention, 832x480 81 frames, one cfg pass"), "wan22-cfg": dict(batch=2, heads=40, tokens=32760, head_dim=128, desc="Wan 2.2 A14B self-attention, 832x480 81 frames, batched cfg"), "ltx2": dict(batch=1, heads=32, tokens=13376, head_dim=128, desc="LTX 2.3 self-attention, 1216x704 121 frames, one cfg pass"), "h3": dict(batch=1, heads=56, tokens=38222, head_dim=128, iters=8, warmup=3, ref_head_chunk=14, sparse=dict(layout=[("text", 0, 512), ("audio", 512, 926), ("video", 926, 38222)]), desc="MiniMax H3 packed self-attention, 1344x768 124 frames (5.2s): 512 text + 414 audio + 37296 video rows, guidance-free"), "h3-long": dict(batch=1, heads=56, tokens=109574, head_dim=128, iters=6, warmup=2, ref_head_chunk=8, config_timeout=1200, sparse=dict(layout=[("text", 0, 512), ("audio", 512, 1718), ("video", 1718, 109574)]), desc="MiniMax H3 packed self-attention, 1344x768 362 frames (15.1s): 512 text + 1206 audio + 107856 video rows"), "masked": dict(batch=1, heads=32, tokens=4608, head_dim=128, mask_fn=key_padding_mask, desc="FLUX.2 Klein shape with boolean key-padding mask, 25% of keys masked"), } # sparse crossover probes at fixed h3 geometry; the smallest token count where a sparse row # beats dense past the verdict threshold is the measured minimum-sequence gate for gate_tokens in (2048, 4096, 8192, 16384, 32768, 65536): shape_presets[f"gate-{gate_tokens // 1024}k"] = dict( batch=1, heads=56, tokens=gate_tokens, head_dim=128, sparse=dict(layout=[("text", 0, 512), ("video", 512, gate_tokens)]), desc=f"sparse crossover probe at h3 geometry, {gate_tokens} tokens", **(dict(iters=8, warmup=3) if gate_tokens >= 32768 else {}), ) full_run = ["sd15", "sdxl", "sdxl-cross", "qwen3-te", "anima", "flux2", "krea2", "wan22", "ltx2", "h3"] sparse_run = ["krea2", "h3", "h3-long"] gate_run = [f"gate-{tokens // 1024}k" for tokens in (2048, 4096, 8192, 16384, 32768, 65536)] # settings advice comes from a self-attention shape with the full config set; cross-attention # and text-encoder shapes measure the hijack's cost there but would mislead as global advice recommendation_presets = ["flux2", "krea2", "anima", "sdxl", "wan22", "ltx2", "h3", "sd15"] default_shapes = "sdxl,flux2" all_sections = ["attention", "dequant", "block"] # combined block benchmark: a dit-style transformer block at krea 2 class geometry, so each row # measures a complete configuration of weights dtype x matmul path x attention end to end # generic dit-block geometries from the model transformer configs: flux.1 (3072 wide, # 24 heads, 4x gelu ff, 4096 image plus 512 text tokens) and krea 2 (6144 wide, 48 heads # after gqa expansion, swiglu at 16384, same joint sequence); optional keys: head_dim # (attention width when heads*head_dim != hidden) and mlp ("gelu" default or "swiglu") block_geometries = { "flux1": dict(hidden=3072, heads=24, mlp_dim=12288, tokens=4608), "krea2": dict(hidden=6144, heads=48, mlp_dim=16384, tokens=4608), # minimax h3: attention wider than the residual stream (56*128 > 5376), swiglu mlp; the # full 124-frame token count makes the block section long, so it runs only when selected "h3": dict(hidden=5376, heads=56, head_dim=128, mlp_dim=14336, mlp="swiglu", tokens=38222, iters=6, warmup=2, config_timeout=900), } default_block_geometries = "flux1,krea2" block_geometry = block_geometries["flux1"] # active geometry; bench_block_section iterates block_attention_specs = { "sdpa": None, # stock torch sdpa "atten int8": dict(matmul_dtype="auto", pv_matmul_dtype="auto"), "atten int8 smooth": dict(matmul_dtype="auto", pv_matmul_dtype="auto", smooth_k=True), "atten int8 hadamard": dict(matmul_dtype="auto", pv_matmul_dtype="auto", use_hadamard=True), "atten full": dict(matmul_dtype="auto", pv_matmul_dtype="int8", smooth_k=True, use_hadamard=True), "atten pv accum": dict(matmul_dtype="auto", pv_matmul_dtype="disabled", use_fp16_accum=True), # the sage-style unsafe mode "atten fp16 accum": dict(matmul_dtype="float16", pv_matmul_dtype="float16", use_fp16_accum=True), # the scaled overflow-proof mode "sage": "sage", # external baselines, resolved to the sage wrappers in build_bench_block "sage fp16 accum": "sagefp16", } # attention-table config id measuring the same kernel as each block attention spec, for the # cross-instrument compute split; specs without a standalone row map to None block_spec_attention_ids = { "sdpa": "base", "atten int8": "int8", "atten int8 smooth": "smooth", "atten int8 hadamard": "hadamard", "atten full": "full", "atten pv accum": "pvaccum", "atten fp16 accum": "fp16full-accum", "sage": "sage", "sage fp16 accum": "sagefp16", } # id, weights config (None = bf16), use quantized matmul, attention spec; fp8/fp4 rows use the # dequant path: quantized matmul auto-selects fp8 for float dtypes, unsupported before sm_89 block_configs = [ ("bf16", None, False, "sdpa"), ("bf16-atten", None, False, "atten int8"), ("bf16-sage", None, False, "sage"), ("bf16-sagefp16", None, False, "sage fp16 accum"), ("int8", dict(weights_dtype="int8"), False, "sdpa"), ("int8-mm", dict(weights_dtype="int8"), True, "sdpa"), ("int8-mm-atten", dict(weights_dtype="int8"), True, "atten int8"), ("int8-mm-smooth", dict(weights_dtype="int8"), True, "atten int8 smooth"), ("int8-mm-hadamard", dict(weights_dtype="int8"), True, "atten int8 hadamard"), ("int8-mm-atten-full", dict(weights_dtype="int8"), True, "atten full"), ("int8-mm-pvaccum", dict(weights_dtype="int8"), True, "atten pv accum"), ("int8-mm-fp16accum", dict(weights_dtype="int8"), True, "atten fp16 accum"), ("int8-mm-sage", dict(weights_dtype="int8"), True, "sage"), ("int8-mm-sagefp16", dict(weights_dtype="int8"), True, "sage fp16 accum"), ("int6", dict(weights_dtype="int6"), False, "sdpa"), ("uint4-mm", dict(weights_dtype="uint4"), True, "sdpa"), ("uint4-mm-atten", dict(weights_dtype="uint4"), True, "atten int8"), ("fp8sdnq", dict(weights_dtype="float8_e4m3fn_sdnq"), False, "sdpa"), ("fp8sdnq-atten", dict(weights_dtype="float8_e4m3fn_sdnq"), False, "atten int8"), ("nvfp4", dict(weights_dtype="float4_e2m1fn", group_size=16), False, "sdpa"), ("nvfp4-atten", dict(weights_dtype="float4_e2m1fn", group_size=16), False, "atten int8"), ] # benchmark configs: id, label, kwargs for sdnq_triton_atten (None = external baseline); # fp8 configs run only on gpus where the float8 probe passes, accum configs only where the # installed sdnq has the use_fp16_accum kwarg bench_configs = [ ("base", "torch sdpa", None), ("sage", "sageattention", None), # label resolved to the dispatched kernel by sage_kernel_label ("sagefp16", "sage int8 qk + fp16 pv, fp16 accum", None), # sm86 only ("amdflash", "triton flash (amd)", None), ("flex", "flex attention, dense", None), # compiled: flex reads its block lists only under compile ("flex-sparse100", "flex + selector, budget 100%", None), # the selector runs but keeps everything, so this row is its overhead alone ("flex-sparse50", "flex + selector, budget 50%", None), ("flex-sparse30", "flex + selector, budget 30%", None), ("flex-sparse15", "flex + selector, budget 15%", None), ("flex-radial30", "flex + static radial band, 30%", None), # density matched control with no per-call producer ("noquant", "sdnq, quantized matmul off", dict(do_quantize=False)), ("int8", "sdnq int8 qk", dict(matmul_dtype="auto", pv_matmul_dtype="auto")), ("int8-sparse100", "sdnq int8 qk + selector, budget 100%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")), # the selector runs but keeps everything, so this row is its overhead on the quantized kernel ("int8-sparse50", "sdnq int8 qk + selector, budget 50%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")), ("int8-sparse30", "sdnq int8 qk + selector, budget 30%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")), ("int8-sparse15", "sdnq int8 qk + selector, budget 15%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")), ("int8-radial30", "sdnq int8 qk + static radial band, 30%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")), # density matched control with no per-call producer ("smooth", "sdnq int8 qk + smooth k", dict(matmul_dtype="auto", pv_matmul_dtype="auto", smooth_k=True)), ("hadamard", "sdnq int8 qk + hadamard", dict(matmul_dtype="auto", pv_matmul_dtype="auto", use_hadamard=True)), ("smooth_hadamard", "sdnq int8 qk + smooth + hadamard", dict(matmul_dtype="auto", pv_matmul_dtype="auto", smooth_k=True, use_hadamard=True)), ("fp16pv", "sdnq int8 qk + fp16 pv", dict(matmul_dtype="auto", pv_matmul_dtype="float16")), ("fp16pv-accum", "sdnq int8 qk + fp16 pv, scaled fp16 accum", dict(matmul_dtype="auto", pv_matmul_dtype="float16", use_fp16_accum=True)), ("int8pv", "sdnq int8 qk + int8 pv", dict(matmul_dtype="auto", pv_matmul_dtype="int8")), ("fp8pv", "sdnq int8 qk + fp8 pv", dict(matmul_dtype="auto", pv_matmul_dtype="float8_e4m3fn")), ("pvaccum", "sdnq int8 qk + unquantized pv, fp16 accum (unsafe)", dict(matmul_dtype="auto", pv_matmul_dtype="disabled", use_fp16_accum=True)), ("full", "sdnq int8 qk + smooth + hadamard + int8 pv", dict(matmul_dtype="auto", pv_matmul_dtype="int8", smooth_k=True, use_hadamard=True)), ("fp16qk", "sdnq fp16 qk", dict(matmul_dtype="float16", pv_matmul_dtype="auto")), ("fp16full", "sdnq fp16 qk + fp16 pv", dict(matmul_dtype="float16", pv_matmul_dtype="float16")), ("fp16full-accum", "sdnq fp16 qk + fp16 pv, scaled fp16 accum", dict(matmul_dtype="float16", pv_matmul_dtype="float16", use_fp16_accum=True)), ("fp8qk", "sdnq fp8 qk", dict(matmul_dtype="float8_e4m3fn", pv_matmul_dtype="auto")), ("fp8full", "sdnq fp8 qk + fp8 pv", dict(matmul_dtype="float8_e4m3fn", pv_matmul_dtype="float8_e4m3fn")), ] # external baselines are compared against but never starred or recommended as sdnq configs; # the unsafe accum mode is measured and displayed under the same rule, since its overflow # tail lives outside what mean error can see # baselines and non-sdnq rows: reported, never recommended as an sdnq setting, and the sparse # rows are lossy by design so a recommendation must not pick one for being fast external_config_ids = ("base", "sage", "sagefp16", "amdflash", "flex") sparse_config_ids = ("flex-sparse100", "flex-sparse50", "flex-sparse30", "flex-sparse15", "flex-radial30", "int8-sparse100", "int8-sparse50", "int8-sparse30", "int8-sparse15", "int8-radial30") unsafe_config_ids = ("pvaccum",) # every preset runs the full config list (availability gates still apply per config); only # hard technical exclusions live here, never runtime trims. sd15: compiling hadamard with # a non pow2 head dim currently hangs torch inductor preset_excluded_configs = {"sd15": {"hadamard", "smooth_hadamard", "full"}} # weight dequant benchmark geometry from the model transformer configs: flux.1 attention # and feed-forward linears (3072 wide), the qwen3 1.7b (anima te) gate/up projection for # text-encoder geometry, and krea 2 12b linears (6144 wide: standalone wq, swiglu gate) dequant_shapes = [ ("flux1 attn 3072x3072", 3072, 3072), ("flux1 mlp 12288x3072", 12288, 3072), ("te mlp 6144x2048", 6144, 2048), ("krea2 wq 6144x6144", 6144, 6144), ("krea2 mlp 16384x6144", 16384, 6144), ] dequant_forward_tokens = 4096 # rows of the forward-bench input, one 1024px image worth of tokens # id, label, sdnq config kwargs; int8 first: the ui default and the dominant pre-quantized format dequant_dtype_configs = [ ("int8", "int8", dict(weights_dtype="int8")), ("uint8", "uint8", dict(weights_dtype="uint8")), ("int6", "int6", dict(weights_dtype="int6")), ("uint6", "uint6", dict(weights_dtype="uint6")), ("int4", "int4", dict(weights_dtype="int4")), ("uint4", "uint4", dict(weights_dtype="uint4")), ("int2", "int2", dict(weights_dtype="int2")), ("uint2", "uint2", dict(weights_dtype="uint2")), ("fp8", "float8_e4m3fn", dict(weights_dtype="float8_e4m3fn")), ("fp8sdnq", "float8_e4m3fn_sdnq", dict(weights_dtype="float8_e4m3fn_sdnq")), ("nvfp4", "float4_e2m1fn g16", dict(weights_dtype="float4_e2m1fn", group_size=16)), ] # svd/hadamard variants measured on top of these dtypes at the first dequant shape; low bits are # where rotation is expected to pay, int8 is the control dequant_variant_dtypes = ["int8", "uint4", "uint2"] dequant_variant_configs = [ ("hadamard", dict(use_hadamard=True)), ("svd", dict(use_svd=True)), ("svd+hadamard", dict(use_svd=True, use_hadamard=True)), ] # float storage dtypes route quantized matmul to fp8 under auto; measure the explicit MatMul # type alternatives so advice on gpus without fp8 matmul cites numbers instead of escape hatches float_mm_dtypes = ["fp8", "fp8sdnq", "nvfp4"] float_mm_alternative_dtypes = ["int8", "float16"] # setting sweeps for the remaining measurable SDNQ options, each at the first dequant shape: # group size 0 = auto, -1 = row-wise; explicit values snap down to a divisor of in_features # and grouped weights force a per-forward re-quantize when quantized matmul is on group_sweep_dtypes = ["int8", "uint4"] group_sweep_values = [0, -1, 32, 64, 128, 256] svd_rank_sweep_dtypes = ["int8", "uint4"] svd_rank_sweep_values = [16, 32, 64, 128] hadamard_group_values = [32, 64, 128, 256] # weight-side rotation group, swept on int8 toggle_dtypes = ["int8", "uint4"] # dequantize using full precision off is measured on these dynamic_quant_requests = ["uint2", "uint4", "int6"] # dynamic quantization escalates until loss passes # conv geometry: sdxl unet mid-block and vae decoder convs, the two conv-heavy model classes conv_shapes = [ ("unet 1280x1280 3x3 @32px", 1280, 1280, 3, 32), ("vae 512x512 3x3 @256px", 512, 512, 3, 256), ] conv_configs = [ # id, weights config (None = bf16 baseline), use conv quantized matmul ("bf16", None, False), ("int8", dict(weights_dtype="int8"), False), ("int8-mm", dict(weights_dtype="int8"), True), ("uint8-mm", dict(weights_dtype="uint8"), True), ] all_dequant_sweeps = ["groups", "svd", "hgroups", "toggles", "conv"] all_mm_backends = ["torch", "triton"] recommend_error_cap = 2.0 # a faster config is not recommended when it multiplies measured output error beyond this recommend_speed_margin = 0.90 # an on/off verdict needs new_ms <= old_ms * this; sub-margin wins sit inside 12-iter run noise and never justify added error atten_settings = [ ("sdnq_attention_matmul_type", "MatMul type"), ("sdnq_attention_pv_matmul_type", "PV MatMul type"), ("sdnq_attention_smooth_k", "Use Smooth K"), ("sdnq_attention_use_hadamard", "Use Hadamard"), ("sdnq_attention_use_fp16_accum", "Use FP16 Accumulation"), ("sdnq_attention_hadamard_group_size", "Hadamard Group Size"), ] def parse_cli(): parser = argparse.ArgumentParser(description="benchmark and validate sdnq attention and weight dequantization on the local gpu") parser.add_argument("--sections", type=str, default=",".join(all_sections), help=f"comma-separated benchmark sections: {', '.join(all_sections)} (default: %(default)s)") parser.add_argument("--dequant-dtypes", type=str, default="all", help=f"comma-separated dequant dtype configs: {', '.join(dtype_id for dtype_id, _label, _cfg in dequant_dtype_configs)}; 'all' runs every one (default: %(default)s)") parser.add_argument("--dequant-variants", type=str, default="all", help=f"comma-separated svd/hadamard variants benched on {', '.join(dequant_variant_dtypes)}: {', '.join(variant_id for variant_id, _cfg in dequant_variant_configs)}; 'all' or 'none' (default: %(default)s)") parser.add_argument("--dequant-sweeps", type=str, default="all", help=f"comma-separated setting sweeps in the dequant section: {', '.join(all_dequant_sweeps)}; 'all' or 'none' (default: %(default)s)") parser.add_argument("--mm-backends", type=str, default="none", help=f"comma-separated quantized-matmul backends to compare in one run: {', '.join(all_mm_backends)}; 'none' benches only the backend this device selects (default: %(default)s)") parser.add_argument("--mm-rounds", type=int, default=2, help="alternating rounds per matmul backend, fastest kept, so clock drift cancels instead of favouring one backend (default: %(default)s)") parser.add_argument("--configs", type=str, default="all", help=f"comma-separated attention configs: {', '.join(config_id for config_id, _label, _kwargs in bench_configs)}; 'all' runs every one (default: %(default)s)") parser.add_argument("--block-configs", type=str, default="all", help=f"comma-separated combined block configs: {', '.join(config_id for config_id, _w, _mm, _a in block_configs)}; 'all' runs every one (default: %(default)s)") parser.add_argument("--block-geometries", type=str, default=default_block_geometries, help=f"comma-separated block geometries: {', '.join(block_geometries)}; 'all' runs every one (default: %(default)s)") parser.add_argument("--shapes", type=str, default=default_shapes, help=f"comma-separated attention shape presets: {', '.join(shape_presets)}; 'all' runs {', '.join(full_run)}, 'sparse' and 'gate' run the sparse and crossover lists (default: %(default)s)") parser.add_argument("--iters", type=int, default=12, help="minimum timed iterations per config, scaled up for fast kernels (default: %(default)s)") parser.add_argument("--warmup", type=int, default=4, help="minimum warmup iterations per config, scaled up for fast kernels (default: %(default)s)") parser.add_argument("--skip-checks", action="store_true", help="skip kernel correctness checks") parser.add_argument("--skip-bench", action="store_true", help="skip benchmarks, run checks and the fp8 and compile probes only") parser.add_argument("--dtype", type=str, default="auto", choices=["auto", "bf16", "fp16"], help="tensor dtype for benchmarks; auto uses the dtype the webui selected for this gpu (default: %(default)s)") parser.add_argument("--config-timeout", type=int, default=None, help="best effort: abort a config whose compile plus first call exceeds this many seconds, 0 disables; cannot interrupt native-level hangs (default: 300, or the limit a preset or block geometry declares for itself)") parser.add_argument("--save", type=str, default="auto", help="plain-text copy of all tables and notes; 'auto' (default) names it -t-.txt in the output directory, 'none' disables, anything else is used as the path") parser.add_argument("--json", type=str, default="auto", help="structured results (environment, probes, per-shape and dequant timings, recommendations); 'auto' (default) names it -t-.json in the output directory, 'none' disables, anything else is used as the path") parser.add_argument("--outdir", type=str, default=None, help="directory for auto-named outputs (default: $SDNQ_BENCH_DIR, or benchmarks/ under the sdnext root)") args = parser.parse_args() args.timeout_flag = args.config_timeout # None lets a preset or block geometry declare its own limit args.config_timeout = resolve_timeout(args.timeout_flag) sys.argv = sys.argv[:1] # sdnext parses argv again on import and rejects unknown arguments return args def package_version(name): try: return importlib.metadata.version(name) except Exception: return None def triton_version(): # distribution name varies by platform (triton-windows, pytorch-triton-rocm); report the module version try: import triton return triton.__version__ except Exception: return None @contextmanager def capture_console_output(): # gate stdout and stderr at the fd level: covers python loggers regardless of their # handler plumbing plus native-code writes (onnxruntime device discovery on wsl). # capture to a file rather than devnull: the sdnext bootstrap calls sys.exit on fatal # startup errors (modules/loader.py, installer.py), so discarding this stream turns a # startup failure into a silent process exit with nothing to debug. # the python-level streams must move with the fds: on a windows console sys.stdout writes # through the console api using the handle behind fd 1, so once fd 1 is a file the next # print raises OSError 'the handle is invalid' and, with stderr equally broken, kills the # interpreter before any handler or finally can run. buffer is deliberately left open: # library loggers built during the import (transformers, torch) capture sys.stderr at # handler construction and would raise on a closed stream long after the window closes captured = {"text": ""} buffer = io.StringIO() saved_stdout_fd, saved_stderr_fd = os.dup(1), os.dup(2) saved_stdout, saved_stderr = sys.stdout, sys.stderr sink_fd, sink_path = tempfile.mkstemp(prefix="sdnq-bench-startup-", suffix=".log") try: sys.stdout.flush() sys.stderr.flush() os.dup2(sink_fd, 1) os.dup2(sink_fd, 2) sys.stdout, sys.stderr = buffer, buffer yield captured finally: os.dup2(saved_stdout_fd, 1) os.dup2(saved_stderr_fd, 2) sys.stdout, sys.stderr = saved_stdout, saved_stderr os.close(sink_fd) os.close(saved_stdout_fd) os.close(saved_stderr_fd) native_text = "" try: with open(sink_path, "r", encoding="utf-8", errors="replace") as fh: native_text = fh.read() os.unlink(sink_path) except OSError: pass captured["text"] = buffer.getvalue() + native_text def load_sdnext(): global shared, devices, sdnq_triton_atten # pylint: disable=global-statement root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if root not in sys.path: sys.path.insert(0, root) if not torch_device_module.is_available(): console.print("[red]no cuda, rocm or xpu device available: sdnq attention requires a gpu with triton[/red]") return False stock_sdpa = torch.nn.functional.scaled_dot_product_attention # the webui startup log (device detect, packages, settings validation) is noise here: the # environment panel reports the stack. the bootstrap reconfigures its loggers during # import and onnxruntime's device discovery warns from c++, so gate the os-level fds for # the import window; on failure the captured log is replayed after the gate lifts os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error") try: with capture_console_output() as startup_log: from modules import shared as shared_module from modules import devices as devices_module from sdnq.kernels.triton_atten import sdnq_triton_atten as atten except BaseException as e: # pylint: disable=broad-exception-caught # SystemExit is not an Exception: the bootstrap exits on a failed torch or library import if isinstance(e, KeyboardInterrupt): raise detail = f"exited with code {e.code}" if isinstance(e, SystemExit) else f"{type(e).__name__}: {e}" console.print(f"[red]sdnext failed to start: {detail}[/red]") text = startup_log["text"].strip() # pylint: disable=used-before-assignment if text: console.print(Panel(escape(text[-4000:]), title="sdnext startup log", box=ROUNDED_BOX)) console.print("run from the sdnext root with the venv active; triton is required") return False # keep the sdnext loggers quiet after the bootstrap too, so stray log lines cannot tear # the live tables mid-bench try: import installer installer.log.setLevel(logging.CRITICAL) from modules.logger import log as sdnext_log sdnext_log.setLevel(logging.CRITICAL) except Exception: pass # importing modules.shared installs the configured sdp override hijacks in this process; # restore stock sdpa so baselines and references measure torch itself torch.nn.functional.scaled_dot_product_attention = stock_sdpa # the triton autotune hijack draws its own rich progress bar on the logger console during # sweeps; a second live display on this tty tramples the benchmark's live tables (stale # border lines left in scrollback), so disarm the bar and keep the hijack's bookkeeping try: from modules import sd_hijack_triton sd_hijack_triton.start_progress = lambda name, total: (None, None) except Exception: pass shared = shared_module devices = devices_module sdnq_triton_atten = atten # inductor dumps failing buffers at CRITICAL to raw stderr, tearing the live # display; failures still surface as exceptions in the tables logging.getLogger("torch._inductor").setLevel(logging.CRITICAL + 1) return True def atten_supports_fp16_accum(): # accum rows need the use_fp16_accum kwarg; skip them on older sdnq builds instead of # failing every row with a TypeError. the kernel entry is wrapped by an inference-context # decorator, so unwrap before reading the signature if sdnq_triton_atten is None: return False try: return "use_fp16_accum" in inspect.signature(inspect.unwrap(sdnq_triton_atten)).parameters except (TypeError, ValueError): return False def atten_supports_block_mask(): # the sdnq sparse rows feed the kernel's block_mask kwarg; skip them on builds without it if sdnq_triton_atten is None: return False try: return "block_mask" in inspect.signature(inspect.unwrap(sdnq_triton_atten)).parameters except (TypeError, ValueError): return False def triton_mm_supports_fp16_accum(): try: from sdnq.kernels import triton_mm, triton_scaled_mm return hasattr(triton_mm, "USE_FP16_ACCUM") and hasattr(triton_scaled_mm, "USE_FP16_ACCUM") except Exception: return False @contextmanager def triton_mm_fp16_accum(): # SDNQ_TRITON_MM_USE_FP16_ACCUM is read once at import into module globals that the # kernel wrappers reread per call (they are triton_op custom ops, so the read happens at # runtime even under torch.compile), and the flag is part of the autotune key, so both # variants cache side by side. sdnq_triton_mm holds its own imported copy of the global, # so both module namespaces get the flip from sdnq.kernels import triton_mm, triton_scaled_mm saved = (triton_mm.USE_FP16_ACCUM, triton_scaled_mm.USE_FP16_ACCUM) triton_mm.USE_FP16_ACCUM = True triton_scaled_mm.USE_FP16_ACCUM = True try: yield finally: triton_mm.USE_FP16_ACCUM, triton_scaled_mm.USE_FP16_ACCUM = saved flex_budgets = {"flex-sparse100": 1.0, "flex-sparse50": 0.50, "flex-sparse30": 0.30, "flex-sparse15": 0.15} sdnq_sparse_budgets = {"int8-sparse100": 1.0, "int8-sparse50": 0.50, "int8-sparse30": 0.30, "int8-sparse15": 0.15} def is_sdnq_sparse(config_id): return config_id in sdnq_sparse_budgets or config_id == "int8-radial30" def make_sdnq_sparse_fn(config_id, q, k, v, attn_mask, kwargs, causal, gqa): """The producer the flex rows time, feeding the quantized kernel's block mask input instead.""" from modules.attention.sparse import selector as sparse_selector def attend(selection): return sdnq_triton_atten(q, k, v, attn_mask=attn_mask, is_causal=causal, enable_gqa=gqa, block_mask=selection.keep, block_mask_m=selection.block_q, block_mask_n=selection.block_kv, **kwargs) if config_id == "int8-radial30": static = sparse_selector.radial_blocks(q.shape[-2], k.shape[-2], 0.30, sparse_selector.SparseSpec(), q.device) return lambda: attend(static) spec = sparse_selector.SparseSpec(budget=sdnq_sparse_budgets[config_id], force=True) cache_key = ("bench", config_id, tuple(q.shape), tuple(k.shape)) return lambda: attend(sparse_selector.select_blocks(q, k, spec, cache_key=cache_key)) def flex_available(): try: import modules.attention.sparse.flex # pylint: disable=unused-import return torch.cuda.is_available() except Exception: return False def make_flex_fn(config_id, q, k, v, scale, gqa): """Time the selector inside the attention it accelerates; a producer measured on its own looks free and is not.""" from modules.attention.sparse import flex as sparse_flex, selector as sparse_selector call = sparse_flex.flex_call() if config_id == "flex": return lambda: call(q, k, v, scale=scale, enable_gqa=gqa) if config_id == "flex-radial30": # a static pattern is built once by construction, which is exactly the advantage it has to defend static = sparse_flex.to_block_mask(sparse_selector.radial_blocks(q.shape[-2], k.shape[-2], 0.30, sparse_selector.SparseSpec(), q.device)) return lambda: call(q, k, v, block_mask=static, scale=scale, enable_gqa=gqa) spec = sparse_selector.SparseSpec(budget=flex_budgets[config_id], force=True) cache_key = ("bench", config_id, tuple(q.shape), tuple(k.shape)) # the webui caches the geometry per layout, so measure that path def run(): selection = sparse_selector.select_blocks(q, k, spec, cache_key=cache_key) return call(q, k, v, block_mask=sparse_flex.to_block_mask(selection), scale=scale, enable_gqa=gqa) return run def sage_attention(): # mirror the backend selection from modules/attention.py: sm86 needs the cuda backend try: if torch_device_module.get_device_capability(torch.device(torch_device)) == (8, 6): from sageattention import sageattn_qk_int8_pv_fp16_cuda def sage_fn(q, k, v, scale): return sageattn_qk_int8_pv_fp16_cuda(q=q, k=k, v=v, tensor_layout="HND", is_causal=False, sm_scale=scale, return_lse=False, pv_accum_dtype="fp32") else: from sageattention import sageattn def sage_fn(q, k, v, scale): return sageattn(q=q, k=k, v=v, attn_mask=None, dropout_p=0.0, is_causal=False, scale=scale) return sage_fn except Exception: return None def sage_kernel_label(): # "sage" is a different kernel per gpu: sageattention's dispatch (core.py sageattn) picks the # pv dtype and accumulator by arch and cuda version, and modules/attention.py forces the cuda # fp16-pv kernel on sm86. name what ran so reports from different gpus stay comparable try: capability = torch_device_module.get_device_capability(torch.device(torch_device)) cuda_version = tuple(int(part) for part in (torch.version.cuda or "0.0").split(".")[:2]) except Exception: return "sageattention" fp8_accum = "fp32+fp16" if cuda_version >= (12, 8) else "fp32+fp32" # sageattention2++ needs cuda 12.8 if capability == (7, 5): return "sage int8 qk + fp16 pv, triton" if capability in {(8, 0), (8, 6), (8, 7)}: return "sage int8 qk + fp16 pv, fp32 accum" if capability == (9, 0): return "sage int8 qk + fp8 pv, fp32+fp32 accum" if capability in {(8, 9), (10, 0), (12, 0), (12, 1)}: return f"sage int8 qk + fp8 pv, {fp8_accum} accum" return "sageattention" def config_label(config_id, label): return sage_kernel_label() if config_id == "sage" else label def amd_triton_flash(): # sdnext's vendored flash attention for rocm/zluda (modules/flash_attn_triton_amd): the # relevant baseline on amd, where stock sdpa has no flash kernel and sage is unavailable try: if getattr(devices, "backend", None) not in {"rocm", "zluda"}: return None from modules.flash_attn_triton_amd import interface_fa def flash_fn(q, k, v, scale, is_causal=False): head_size = q.size(3) if head_size % 8 != 0: pad = 8 - head_size % 8 q = torch.nn.functional.pad(q, [0, pad]) k = torch.nn.functional.pad(k, [0, pad]) v = torch.nn.functional.pad(v, [0, pad]) q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) out = torch.zeros_like(q) interface_fa.fwd(q, k, v, out, 0.0, scale, is_causal) return out[..., :head_size].transpose(1, 2) return flash_fn except Exception: return None def sage_attention_fp16_accum(): # sage's fast sm86 mode: fp16 pv accumulation, which sdnext deliberately does not ship # (accumulator overflow risk on extreme activations); benched to quantify what that # choice costs in speed and buys in error try: if torch_device_module.get_device_capability(torch.device(torch_device)) != (8, 6): return None from sageattention import sageattn_qk_int8_pv_fp16_cuda def sage_fn(q, k, v, scale): return sageattn_qk_int8_pv_fp16_cuda(q=q, k=k, v=v, tensor_layout="HND", is_causal=False, sm_scale=scale, return_lse=False, pv_accum_dtype="fp16") return sage_fn except Exception: return None def make_qkv(batch, heads, tokens, head_dim, structured=True, kv_heads=None, kv_tokens=None): # structured keys carry a shared per-channel bias plus a few outlier channels, mimicking the # key statistics that motivate smoothing and rotation; absolute error varies by model # architecture while the relative ordering of configs holds generator = torch.Generator(device=torch_device).manual_seed(1234) kv_heads = kv_heads or heads kv_tokens = kv_tokens or tokens q = torch.randn(batch, heads, tokens, head_dim, device=torch_device, dtype=bench_dtype, generator=generator) k = torch.randn(batch, kv_heads, kv_tokens, head_dim, device=torch_device, dtype=bench_dtype, generator=generator) v = torch.randn(batch, kv_heads, kv_tokens, head_dim, device=torch_device, dtype=bench_dtype, generator=generator) if structured: k = k + torch.randn(1, kv_heads, 1, head_dim, device=torch_device, dtype=bench_dtype, generator=generator) * 3.0 k[..., [3, head_dim // 2, head_dim - 5]] *= 4.0 return q, k, v def fp32_reference(q, k, v, head_chunk=0, **kwargs): # sdnext enables tf32 globally; a math-backend dispatch fallback would degrade the reference to tf32 precision tf32_matmul = torch.backends.cuda.matmul.allow_tf32 tf32_cudnn = torch.backends.cudnn.allow_tf32 torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False try: if head_chunk and not kwargs.get("enable_gqa") and q.shape[1] > head_chunk: # head-sliced reference: bounds the fp32 peak on very long sequences; gqa shapes # keep the one-shot path since slicing q heads would have to regroup kv heads attn_mask = kwargs.pop("attn_mask", None) outs = [] for start in range(0, q.shape[1], head_chunk): heads = slice(start, start + head_chunk) mask_slice = attn_mask[:, heads] if attn_mask is not None and attn_mask.shape[1] > 1 else attn_mask outs.append(torch.nn.functional.scaled_dot_product_attention(q[:, heads].to(torch.float32), k[:, heads].to(torch.float32), v[:, heads].to(torch.float32), attn_mask=mask_slice, **kwargs)) return torch.cat(outs, dim=1) return torch.nn.functional.scaled_dot_product_attention(q.to(torch.float32), k.to(torch.float32), v.to(torch.float32), **kwargs) finally: torch.backends.cuda.matmul.allow_tf32 = tf32_matmul torch.backends.cudnn.allow_tf32 = tf32_cudnn def rel_err(out, ref): out = out.to(torch.float32) return ((out - ref).norm() / ref.norm()).item() def max_token_err(out, ref, eps=1e-6): # worst single-token relative error over the feature axis: catches localized corruption # (one garbage image region, one dead token) that a global norm averages away out = out.to(torch.float32) numerator = (out - ref).norm(dim=-1) denominator = ref.norm(dim=-1).clamp_min(eps) return (numerator / denominator).max().item() def err_cell(value, fmt="{:.5f}"): # non-finite outputs are the black-image failure class; label them instead of printing nan if value is None: return "-" if not math.isfinite(value): return "[red]non-finite[/red]" return fmt.format(value) def fp32_linear_reference(x, weight_fp32): # sdnext enables tf32 globally; compute the linear reference in true fp32, matching fp32_reference tf32_matmul = torch.backends.cuda.matmul.allow_tf32 torch.backends.cuda.matmul.allow_tf32 = False try: return torch.nn.functional.linear(x.to(torch.float32), weight_fp32) finally: torch.backends.cuda.matmul.allow_tf32 = tf32_matmul def cuda_context_alive(): # a kernel fault (misaligned address, illegal memory access) is sticky: every later cuda # call in the process fails; detect it so sections abort cleanly instead of cascading try: torch.zeros(1, device=torch_device) torch_device_module.synchronize() return True except Exception: return False def error_summary(e, limit=60): # dynamo appends advice lines ("Set TORCHDYNAMO_VERBOSE=1 ...") after the real error; # keep the last substantive line so table rows show the actual failure lines = [line.strip() for line in str(e).splitlines()] lines = [line for line in lines if line and "TORCHDYNAMO_VERBOSE" not in line and "TORCH_LOGS" not in line] text = lines[-1] if lines else type(e).__name__ return text[:limit] compile_error_markers = ("InductorError", "BackendCompilerFailed", "CompilationError", "TritonError") def failure_text(e): # label known compile failures plainly, technical detail dimmed detail = error_summary(e) if isinstance(e, TimeoutError): return f"[red]failed: compile timed out[/red] [dim]({detail})[/dim]" if type(e).__name__ in compile_error_markers or any(marker in str(e) for marker in compile_error_markers): return f"[red]failed: torch compile error[/red] [dim]({detail})[/dim]" return f"[red]{type(e).__name__}[/red]: {detail}" def speedup_cell(base_ms, ms): ratio = base_ms / ms text = f"x{ratio:4.2f}" if ratio >= 1.10: return f"[green]{text}[/green]" if ratio < 0.95: return f"[red]{text}[/red]" return text def live_progress(): # status line shown under a live-filling table progress = Progress(SpinnerColumn(), TextColumn("{task.description}"), TimeElapsedColumn(), console=console) task = progress.add_task("starting", total=None) return progress, task def resolve_timeout(flag, declared=None): # the cli flag wins when given; otherwise a preset or block geometry may declare its own limit if flag is not None: return flag return 300 if declared is None else declared @contextmanager def time_limit(seconds, label): # torch.compile can spin indefinitely in sympy/inductor on pathological graphs # (e.g. hadamard over a padded non pow2 head dim); turn a wedged compile into a red row if seconds <= 0 or not hasattr(signal, "SIGALRM"): yield return def handler(signum, frame): # pylint: disable=unused-argument raise TimeoutError(f"{label}: compile plus first call exceeded {seconds}s") previous = signal.signal(signal.SIGALRM, handler) signal.alarm(seconds) try: yield finally: signal.alarm(0) signal.signal(signal.SIGALRM, previous) def bench_stats(fn, warmup, iters, on_phase=None): """Time fn and return (median ms, relative sigma of that median). The sigma covers within-sample dispersion only (IQR-based standard error of a median); slow clock and thermal drift across a run is tracked separately by the per-shape sentinel re-measurements and folded in at verdict time.""" if on_phase: on_phase("warmup") fn() torch_device_module.synchronize() started = time.perf_counter() fn() torch_device_module.synchronize() estimate = max(time.perf_counter() - started, 1e-6) # scale counts to a time budget so short kernels get enough activity to ramp gpu clocks warmup = min(max(warmup, int(0.2 / estimate)), 500) iters = min(max(iters, int(0.5 / estimate)), 500) for _ in range(warmup): fn() torch_device_module.synchronize() if on_phase: on_phase(f"timing {iters} iterations") events = [(torch_device_module.Event(enable_timing=True), torch_device_module.Event(enable_timing=True)) for _ in range(iters)] for start, end in events: start.record() fn() end.record() torch_device_module.synchronize() times = sorted(start.elapsed_time(end) for start, end in events) n = len(times) median = times[n // 2] iqr = times[(3 * n) // 4] - times[n // 4] # standard error of a median: 1.2533 * sigma / sqrt(n), robust sigma from IQR / 1.349 sigma_rel = (1.2533 * (iqr / 1.349) / math.sqrt(n)) / median if median > 0 else 0.0 return median, sigma_rel def bench(fn, warmup, iters, on_phase=None): return bench_stats(fn, warmup, iters, on_phase)[0] drift_samples = [] # |ln ratio| of per-shape sentinel re-measurements; tracks run-level clock drift def record_drift(first_ms, repeat_ms): if first_ms and repeat_ms: drift_samples.append(abs(math.log(repeat_ms / first_ms))) def run_drift_sigma(): # rms of the sentinel deltas, floored so a lucky pair of samples cannot claim a # noiseless run; replays of json files inject the stored value via drift_override if drift_override is not None: return drift_override if not drift_samples: return 0.0 rms = math.sqrt(sum(d * d for d in drift_samples) / len(drift_samples)) return min(max(rms, 0.005), 0.10) drift_override = None verdict_z = 1.28 # one-sided 90%: an on/off verdict is only stated when its margin test clears this # split instruments A and B differ systematically (strided views out of the fused projection vs # contiguous standalone tensors): 0.1-2.5% over 19 same-kernel rows at h3, far below the 17-25% hadamard gap split_instrument_offset = 0.05 def sidak_z_for(count): # per-test z holding the family-wise confidence of verdict_z across count candidates # (sidak): the selected best of several noisy rows sits low by selection, so testing # it at the single-test z would overstate confidence if count <= 1: return verdict_z single = statistics.NormalDist().cdf(verdict_z) return statistics.NormalDist().inv_cdf(single ** (1.0 / count)) # the synthetic-tensor errors are seed-fixed math and land within a few percent across # nvidia, amd and intel gpus; a value outside its band means the measurement itself is # off (wrong dtype, broken kernel, degraded reference), not an interesting gpu error_sanity_bands = { # (shape, config) -> (low, high), bfloat16 runs only ("anima", "smooth"): (0.0159, 0.0215), ("sdxl", "smooth"): (0.0185, 0.0250), ("anima", "noquant"): (0.0021, 0.0043), ("sdxl", "noquant"): (0.0025, 0.0050), } def row_sigma(entry, key="ms"): # relative sigma of one measured row: within-sample dispersion plus the run-level # drift sampled by the sentinels; None when the row predates sigma capture, which # drops the verdict back to a point-estimate test (old json replays) boot = entry.get(f"{key}_sigma") if isinstance(entry, dict) else None drift = run_drift_sigma() if boot is None: return drift or None return math.sqrt(boot * boot + drift * drift) def pair_sigma(entry_a, key_a, entry_b, key_b): sa, sb = row_sigma(entry_a, key_a), row_sigma(entry_b, key_b) if sa is None and sb is None: return None return math.sqrt((sa or 0.0) ** 2 + (sb or 0.0) ** 2) def speed_verdict(new_ms, old_ms, sigma=None, margin=None, z=verdict_z): """Three-zone test of "new beats old by the speed margin": returns 'faster', 'not_faster', or 'inconclusive' when the gap sits inside z*sigma of the margin. Without a sigma the test degrades to the plain point-estimate threshold.""" if not (new_ms and old_ms): return None margin = margin if margin is not None else recommend_speed_margin gap = math.log(old_ms / new_ms) + math.log(margin) # > 0 clears the margin if sigma: if gap > z * sigma: return "faster" if gap < -z * sigma: return "not_faster" return "inconclusive" return "faster" if gap > 0 else "not_faster" def free_vram_gb(): free, _total = torch_device_module.mem_get_info() return free / 1024**3 def fp8_compile_gate_flag(): # False on gpus where sdnq upcasts e4m3 storage to the scale dtype before the compiled # dequant, because triton cannot convert e4m3 there; absent on builds without the gate try: from sdnq import kernel_wrappers as sdnq_kernel_wrappers return getattr(sdnq_kernel_wrappers, "is_fp8_compile_supported", None) except Exception: return None def fp8_failure_is_capability(detail): # the triton pre-sm_89 signature; anything else is an environment failure where a torch # or triton issue is more likely than a missing hardware capability detail = detail or "" return "fp8e4nv" in detail or "not supported in this architecture" in detail def print_environment(fp8_result, prep_status, prep_detail, weight_dequant_result=None): device = torch.device(torch_device) capability = torch_device_module.get_device_capability(device) # backend runtime versions (cuda/cudnn/driver, hip, ipex, openvino) so a shared # report identifies the stack without inferring it from the torch version string try: gpu_info = devices.get_gpu_info() or {} except Exception: gpu_info = {} runtime_versions = {key: gpu_info[key] for key in ("cuda", "hip", "cudnn", "driver", "ipex", "openvino") if gpu_info.get(key)} runtime_line = f"python: {sys.version.split()[0]} ({sys.platform}) backend: {getattr(devices, 'backend', 'unknown')}" if runtime_versions: runtime_line += " " + " ".join(f"{key}: {value}" for key, value in runtime_versions.items()) lines = [ f"device: [cyan]{torch_device_module.get_device_name(device)}[/cyan] capability={capability}", f"torch: {torch.__version__} triton: {triton_version() or 'not installed'} sageattention: {package_version('sageattention') or 'not installed'} flash-attn: {package_version('flash-attn') or 'not installed'}", runtime_line, f"benchmark dtype: {dtype_label()}", ] if fp8_result is not None: if fp8_result["qk"][0]: lines.append("float8_e4m3fn matmul: [green]supported[/green]") elif fp8_failure_is_capability(fp8_result["qk"][1]): lines.append(f"float8_e4m3fn matmul: [red]not supported on this gpu, selecting it fails generation[/red] [dim]({escape(fp8_result['qk'][1])})[/dim]") else: lines.append(f"float8_e4m3fn matmul: [red]failed to compile in this environment, selecting it fails generation[/red]; the error is not the hardware-capability signature, a torch or triton issue is more likely than the gpu [dim]({escape(fp8_result['qk'][1])})[/dim]") lines.append(f"sdnq attention enabled in current config: {'[green]yes[/green]' if 'SDNQ attention' in shared.opts.cross_attention_optimization else '[yellow]no, enable via Compute Settings -> Cross Attention (requires restart)[/yellow]'}") if prep_status == "disabled": lines.append("compiled input prep: torch.compile disabled in config, input prep runs eager") elif prep_status == "working": lines.append("compiled input prep: [green]working[/green]") else: lines.append(f"compiled input prep: [red]failing, every sdnq attention call errors at generation[/red] [dim]({escape(prep_detail)})[/dim]") if prep_status == "failing_dynamic": lines.append(" fix, verified on this machine: set SDNQ_COMPILE_KWARGS='{\"dynamic\": false}' (recompiles per shape); or install msvc build tools; or set SDNQ_USE_TORCH_COMPILE=0") else: lines.append(" fix: install a host c++ compiler (msvc build tools on windows), or set SDNQ_USE_TORCH_COMPILE=0") if weight_dequant_result is not None: gate_flag = fp8_compile_gate_flag() e4m3_ok, e4m3_detail = weight_dequant_result["float8_e4m3fn"] e5m2_ok, e5m2_detail = weight_dequant_result["float8_e5m2"] if e4m3_ok: lines.append("compiled weight dequant, float8_e4m3fn storage: [green]supported[/green]") else: lines.append(f"compiled weight dequant, raw float8_e4m3fn storage: [red]fails on this gpu[/red] [dim]({escape(e4m3_detail)})[/dim]") if gate_flag is False: lines.append(f" sdnq upcasts e4m3 weights to {dtype_label()} before the compiled dequant, so generation is safe and the fp8 rows below are measured that way (SDNQ_ALLOW_FP8_COMPILE overrides)") elif gate_flag is None: lines.append(" [red]this sdnq build has no fp8 compile gate: loading fp8 storage weights with sdnq torch.compile active fails at generation (SDNQ_USE_TORCH_COMPILE=0 avoids it)[/red]") e5m2_verdict = "[green]supported[/green]" if e5m2_ok else f"[red]fails[/red] [dim]({escape(e5m2_detail)})[/dim]" lines.append(f"compiled weight dequant, float8_e5m2 storage: {e5m2_verdict}") if not atten_supports_fp16_accum(): lines.append("fp16 accumulation kwarg: [yellow]absent in this sdnq build, accum rows skipped[/yellow]") if not atten_supports_block_mask(): lines.append("block mask kwarg: [yellow]absent in this sdnq build, sdnq sparse rows skipped[/yellow]") overrides = [f"{key}={value}" for key, value in os.environ.items() if key.startswith("SDNQ_TRITON_ATTEN") or key.startswith("SDNQ_TRITON_MM") or key.startswith("SDNQ_ALLOW_FP8") or key.startswith("SDNQ_COMPILE")] if overrides: lines.append(f"env overrides: {' '.join(overrides)}") emit(Panel("\n".join(lines), title="environment", box=ROUNDED_BOX)) report["environment"] = dict( device=torch_device_module.get_device_name(device), capability=f"{capability[0]}.{capability[1]}", python=sys.version.split()[0], platform=sys.platform, backend=str(getattr(devices, "backend", None)), torch=str(torch.__version__), triton=triton_version(), sageattention=package_version("sageattention"), dtype=dtype_label(), **runtime_versions, fp8_attention_matmul=fp8_result["qk"][0] if fp8_result is not None else None, atten_fp16_accum=atten_supports_fp16_accum(), atten_block_mask=atten_supports_block_mask(), triton_mm_fp16_accum=os.environ.get("SDNQ_TRITON_MM_USE_FP16_ACCUM", None), compiled_input_prep=prep_status, fp8_compile_gate=fp8_compile_gate_flag(), weight_dequant_compile={name: dict(ok=ok, detail=detail) for name, (ok, detail) in (weight_dequant_result or {}).items()}, ) def resolve_output_paths(args): # auto outputs follow the archive convention -t-, one # stem shared by the txt/json pair, suffixed -2, -3... when the pair already exists args.save = None if str(args.save).strip().lower() in {"none", ""} else args.save args.json = None if str(args.json).strip().lower() in {"none", ""} else args.json if args.save != "auto" and args.json != "auto": return try: gpu = torch_device_module.get_device_name(torch.device(torch_device)) except Exception: gpu = "gpu" drop = {"nvidia", "geforce", "rtx", "gtx", "amd", "radeon", "rx", "intel", "arc", "graphics", "apple", "laptop", "gpu"} tokens = ["".join(ch for ch in token if ch.isalnum()) for token in gpu.lower().replace("(r)", " ").replace("(tm)", " ").split()] tokens = [token for token in tokens if token and token not in drop] gpu_id = "".join(tokens) or "gpu" torch_digits = "".join(str(torch.__version__).split("+", maxsplit=1)[0].split(".")[:2]) stem = f"{gpu_id}-t{torch_digits}-{time.strftime('%b%d').lower()}" outdir = args.outdir or os.environ.get("SDNQ_BENCH_DIR", None) or os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "benchmarks") outdir = os.path.expanduser(outdir) os.makedirs(outdir, exist_ok=True) candidate, counter = stem, 1 while os.path.exists(os.path.join(outdir, f"{candidate}.txt")) or os.path.exists(os.path.join(outdir, f"{candidate}.json")): counter += 1 candidate = f"{stem}-{counter}" if args.save == "auto": args.save = os.path.join(outdir, f"{candidate}.txt") if args.json == "auto": args.json = os.path.join(outdir, f"{candidate}.json") def print_banner(selected, sections, args): lines = [ "measures sdnq attention and weight dequantization speed and accuracy on this gpu and recommends values for [cyan]Compute Settings -> SDNQ / SDNQ Attention[/cyan]", f"sections: [cyan]{', '.join(sections)}[/cyan]", ] if args.skip_bench: lines.append("running correctness checks and the float8 and compile probes only (--skip-bench)") elif "attention" in sections: lines.append(f"attention shapes: [cyan]{', '.join(selected)}[/cyan] (available: {', '.join(shape_presets)}; pass --shapes to match the models you use)") lines.append("first run compiles triton kernels per shape and can take several minutes; repeat runs are much faster") if args.mm_backends.strip().lower() not in {"none", ""} and "dequant" in sections: lines.append(f"matmul backends compared in one run: [cyan]{args.mm_backends}[/cyan] (swapped in-process, {args.mm_rounds} alternating rounds per dtype)") if args.save: lines.append(f"a plain-text copy of the results will be saved to [cyan]{args.save}[/cyan]") if args.json: lines.append(f"structured results will be saved to [cyan]{args.json}[/cyan]") console.print(Panel("\n".join(lines), title="sdnq benchmark", box=ROUNDED_BOX)) def probe_fp8(): # hardware probe, prep runs eager via the module-global swap so the verdict isolates the # attention kernel: fp8 kernel compile fails on pre-ada nvidia and pre-rdna4/cdna3 amd; # prep failures are dtype-independent and probed separately. do not force eager by # toggling torch._dynamo.config.disable: torch 2.13+ raises "found no compiled frames" # for fullgraph-compiled functions called inside a disable window try: from sdnq import kernel_wrappers as sdnq_kernel_wrappers is_fp8_mm_supported = getattr(sdnq_kernel_wrappers, "is_fp8_mm_supported", True) except Exception: is_fp8_mm_supported = True if not is_fp8_mm_supported: return dict(qk=(False, "FP8 matmul is not supported in this architecture"), pv=(False, "FP8 matmul is not supported in this architecture")) from sdnq.kernels import triton_atten as atten_module q, k, v = make_qkv(1, 2, 256, 64, structured=False) result = {} compiled_prep = atten_module.get_attn_inputs inner_prep = getattr(compiled_prep, "_torchdynamo_orig_callable", None) if inner_prep is not None: atten_module.get_attn_inputs = inner_prep try: with console.status("probing float8 support") as status: for name, kwargs in [("qk", dict(matmul_dtype="float8_e4m3fn", pv_matmul_dtype="auto")), ("pv", dict(matmul_dtype="auto", pv_matmul_dtype="float8_e4m3fn"))]: status.update(f"probing float8 support: {name} matmul (a compile failure here is normal on older gpus)") try: out = sdnq_triton_atten(q, k, v, **kwargs) torch_device_module.synchronize() result[name] = (not bool(torch.isnan(out).any().item()), "compiles and runs") except Exception as e: result[name] = (False, f"{type(e).__name__}: {error_summary(e, 120)}") finally: atten_module.get_attn_inputs = compiled_prep return result def probe_compiled_prep(): # inductor lowers part of the dynamic-shape prep to a cpu helper kernel, so a missing # host c++ compiler (msvc on windows) fails every sdnq attention call at generation. # must run before any other sdnq_triton_atten call: a prior eager run masks the # cold-start failure the webui hits from sdnq.common import use_torch_compile if not use_torch_compile: return "disabled", None q, k, v = make_qkv(1, 2, 256, 64, structured=False) with console.status("probing compiled input prep (compiles on first run, cached afterwards)"): try: sdnq_triton_atten(q, k, v, matmul_dtype="auto", pv_matmul_dtype="auto") torch_device_module.synchronize() return "working", None except Exception as e: torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state detail = f"{type(e).__name__}: {error_summary(e, 120)}" # check whether the dynamic=false workaround holds from sdnq.kernels import triton_atten as atten_module compiled_prep = atten_module.get_attn_inputs inner = getattr(compiled_prep, "_torchdynamo_orig_callable", None) if inner is None: return "failing", detail static_ok = False with console.status("compiled input prep failed, verifying the dynamic=false workaround"): atten_module.get_attn_inputs = torch.compile(inner, fullgraph=True, dynamic=False) try: sdnq_triton_atten(q, k, v, matmul_dtype="auto", pv_matmul_dtype="auto") torch_device_module.synchronize() static_ok = True except Exception: pass finally: atten_module.get_attn_inputs = compiled_prep torch._dynamo.reset() # pylint: disable=protected-access return ("failing_dynamic" if static_ok else "failing"), detail def make_source_weight(out_features, in_features, seed=1234): # structured weights with a few high-magnitude input channels, mimicking the outlier # statistics of real dit linears; keeps quantization error columns in a realistic range generator = torch.Generator(device=torch_device).manual_seed(seed) weight = torch.randn(out_features, in_features, device=torch_device, dtype=bench_dtype, generator=generator) * 0.02 weight[:, [3, in_features // 2, in_features - 5]] *= 8.0 return weight class BenchBlock(torch.nn.Module): # dit-style block: fused qkv self-attention plus a gelu or swiglu mlp, both with residuals; # head_dim decouples attention width from hidden for models whose attention is wider than # the residual stream; the attention_fn attribute is set per benchmark config def __init__(self, hidden, heads, mlp_dim, head_dim=None, mlp="gelu", device=None, dtype=None): super().__init__() self.heads = heads self.head_dim = head_dim if head_dim is not None else hidden // heads self.mlp = mlp inner = heads * self.head_dim self.norm1 = torch.nn.LayerNorm(hidden, elementwise_affine=False, device=device, dtype=dtype) self.norm2 = torch.nn.LayerNorm(hidden, elementwise_affine=False, device=device, dtype=dtype) self.qkv = torch.nn.Linear(hidden, inner * 3, bias=False, device=device, dtype=dtype) self.proj = torch.nn.Linear(inner, hidden, bias=False, device=device, dtype=dtype) self.up = torch.nn.Linear(hidden, mlp_dim, bias=False, device=device, dtype=dtype) if mlp == "swiglu": self.gate = torch.nn.Linear(hidden, mlp_dim, bias=False, device=device, dtype=dtype) self.down = torch.nn.Linear(mlp_dim, hidden, bias=False, device=device, dtype=dtype) self.attention_fn = None def forward(self, x): batch, tokens, _channels = x.shape h = self.norm1(x) qkv = self.qkv(h).view(batch, tokens, 3, self.heads, self.head_dim).permute(2, 0, 3, 1, 4) attn = self.attention_fn(qkv[0], qkv[1], qkv[2]).transpose(1, 2).reshape(batch, tokens, self.heads * self.head_dim) x = x + self.proj(attn) h = self.norm2(x) if self.mlp == "swiglu": return x + self.down(torch.nn.functional.silu(self.gate(h)) * self.up(h)) return x + self.down(torch.nn.functional.gelu(self.up(h))) def build_block_module(dtype=None): # construct a block for the active geometry; every construction site goes through here so # geometry keys are read in exactly one place return BenchBlock( block_geometry["hidden"], block_geometry["heads"], block_geometry["mlp_dim"], head_dim=block_geometry.get("head_dim"), mlp=block_geometry.get("mlp", "gelu"), device=torch_device, dtype=dtype if dtype is not None else bench_dtype, ) def make_block_master(): # one master weight set shared by every block config, so all rows quantize identical weights; # seed order keeps gelu geometries bitwise stable, swiglu appends its gate after up block = build_block_module() linears = [block.qkv, block.proj, block.up, block.down] if hasattr(block, "gate"): linears.append(block.gate) with torch.no_grad(): for seed, linear in enumerate(linears, start=1): linear.weight.copy_(make_source_weight(linear.out_features, linear.in_features, seed=seed)) return {key: value.clone() for key, value in block.state_dict().items()} def build_bench_block(master_sd, weights_cfg, use_mm, attention_spec): from sdnq import SDNQConfig from sdnq.quantizer import apply_sdnq_to_module block = build_block_module() block.load_state_dict(master_sd) block.eval() for param in block.parameters(): param.requires_grad_(False) if weights_cfg is not None: config = SDNQConfig(use_quantized_matmul=use_mm, add_skip_keys=False, **weights_cfg) block, _config = apply_sdnq_to_module(block, config, torch_dtype=bench_dtype) atten_kwargs = block_attention_specs[attention_spec] if atten_kwargs is None: def attention_fn(q, k, v): return torch.nn.functional.scaled_dot_product_attention(q, k, v) elif isinstance(atten_kwargs, str): sage_fn = sage_attention() if atten_kwargs == "sage" else sage_attention_fp16_accum() if sage_fn is None: raise RuntimeError("sageattention unavailable") def attention_fn(q, k, v, fn=sage_fn): return fn(q, k, v, q.shape[-1] ** -0.5) else: def attention_fn(q, k, v, kw=atten_kwargs): return sdnq_triton_atten(q, k, v, **kw) block.attention_fn = attention_fn return block def block_storage_bytes(block): total = 0 for module in block.modules(): if hasattr(module, "sdnq_dequantizer"): total += layer_storage_bytes(module) elif isinstance(module, torch.nn.Linear): total += module.weight.numel() * module.weight.element_size() return total def make_quantized_linear(weight, weights_dtype, group_size=0, use_quantized_matmul=False, use_svd=False, use_hadamard=False, quantized_matmul_dtype=None, device=torch_device, **config_kwargs): # quantize through the same entry point model loading uses, so forward benches measure # the production wrapper classes and dequantizer configuration; returns the quantize wall # time, a one-shot measurement of what on-the-fly quantization pays per layer at load from sdnq import SDNQConfig from sdnq.quantizer import sdnq_quantize_layer out_features, in_features = weight.shape linear = torch.nn.Linear(in_features, out_features, bias=False, device=device, dtype=bench_dtype) with torch.no_grad(): linear.weight.copy_(weight) config = SDNQConfig(weights_dtype=weights_dtype, group_size=group_size, use_quantized_matmul=use_quantized_matmul, use_svd=use_svd, use_hadamard=use_hadamard, quantized_matmul_dtype=quantized_matmul_dtype, add_skip_keys=False, **config_kwargs) torch_device_module.synchronize() started = time.perf_counter() layer, _config = sdnq_quantize_layer(linear, config, torch_dtype=bench_dtype, param_name="bench.weight") torch_device_module.synchronize() quant_seconds = time.perf_counter() - started if not hasattr(layer, "sdnq_dequantizer"): raise RuntimeError(f"sdnq did not quantize the layer to {weights_dtype}") return layer, quant_seconds def layer_storage_bytes(layer): # measured storage of everything the quantized layer keeps: packed weights, scales, # zero points and svd factors; sizes on disk and in vram follow this, not the nominal bits total = 0 for name in ("weight", "scale", "zero_point", "svd_up", "svd_down"): tensor = getattr(layer, name, None) if isinstance(tensor, (torch.Tensor, torch.nn.Parameter)): total += tensor.numel() * tensor.element_size() return total def size_cell(size_bytes): return f"{size_bytes / 1024**2:6.1f}mb" def quant_cell(quant_seconds): return f"{quant_seconds * 1000.0:5.0f}ms" if quant_seconds is not None else "-" def dequant_args(layer, upcast_fp8=True): # mirror SDNQDequantizer.__call__'s marshaling so the direct compiled call below measures # the same graph the instance would compile, including the e4m3 upcast sdnq applies before # the compiled dequant on gpus where triton cannot convert it (pre sm_89). # upcast_fp8=False keeps the hardware probe honest: it must compile the raw e4m3 weight deq = layer.sdnq_dequantizer weight = layer.weight if upcast_fp8 and weight.dtype == torch.float8_e4m3fn and fp8_compile_gate_flag() is False: weight = weight.to(dtype=layer.scale.dtype) kwargs = dict( zero_point=getattr(layer, "zero_point", None), svd_up=getattr(layer, "svd_up", None), svd_down=getattr(layer, "svd_down", None), dtype=deq.result_dtype, result_shape=deq.result_shape, quantized_weight_shape=deq.quantized_weight_shape, re_quantize_for_matmul=deq.re_quantize_for_matmul or deq.is_packed, ) return (deq.weights_dtype, weight, layer.scale), kwargs def get_compiled_dequantize_weight(): # sdnq's own dequantize_weight_compiled is a passthrough when the config has Dequantize # using torch.compile off; compile a tool-owned variant with the same kwargs so compiled # rows always measure what the webui runs with the option on. sdnq only raises the dynamo # recompile limits when the option is on; raise them here too or the per-dtype-per-shape # specializations exceed the default limit of 8 and silently fall back to eager global compiled_dequantize_weight # pylint: disable=global-statement if compiled_dequantize_weight is None: for limit_name in ("recompile_limit", "cache_size_limit", "accumulated_recompile_limit", "accumulated_cache_size_limit"): if hasattr(torch._dynamo.config, limit_name): # pylint: disable=protected-access setattr(torch._dynamo.config, limit_name, max(8192, getattr(torch._dynamo.config, limit_name) or 0)) # pylint: disable=protected-access from sdnq.dequantizer import dequantize_weight compiled_dequantize_weight = torch.compile(dequantize_weight, fullgraph=True, dynamic=False) return compiled_dequantize_weight def probe_weight_dequant_compile(): # compiled weight dequantization per fp8 storage dtype: triton before sm_89 lacks e4m3 # conversions, so compiled dequant of raw float8_e4m3fn storage crashes on ampere while # e5m2 is expected to compile; the verdicts show whether sdnq's e4m3 upcast is load-bearing # on this gpu, so probe the raw weight rather than the upcast one the bench rows use result = {} for name in ("float8_e4m3fn", "float8_e5m2"): with console.status(f"probing compiled weight dequant: {name} storage (a compile failure here is expected for e4m3 on pre-ada nvidia)"): try: layer, _quant_seconds = make_quantized_linear(make_source_weight(256, 256), name) args, kwargs = dequant_args(layer, upcast_fp8=False) out = get_compiled_dequantize_weight()(*args, **kwargs) torch_device_module.synchronize() result[name] = (bool(torch.isfinite(out).all().item()), "compiles and runs") except Exception as e: result[name] = (False, f"{type(e).__name__}: {error_summary(e, 120)}") reset_compiled_dequant() # drop failed compile state so later rows compile fresh return result def reset_compiled_dequant(): global compiled_dequantize_weight # pylint: disable=global-statement compiled_dequantize_weight = None torch._dynamo.reset() # pylint: disable=protected-access def run_correctness(): checks = [] q, k, v = make_qkv(2, 8, 512, 64, structured=False) checks.append(("plain", (q, k, v), {}, {})) checks.append(("causal", (q, k, v), dict(is_causal=True), {})) qg, kg, vg = make_qkv(2, 8, 512, 64, structured=False, kv_heads=2) checks.append(("gqa 8:2 heads", (qg, kg, vg), dict(enable_gqa=True), {})) bool_mask = torch.zeros(2, 1, 1, 512, device=torch_device, dtype=torch.bool) bool_mask[..., :384] = True checks.append(("bool key-padding mask", (q, k, v), dict(attn_mask=bool_mask), {})) float_mask = torch.zeros(2, 1, 1, 512, device=torch_device, dtype=bench_dtype) float_mask[..., 384:] = float("-inf") checks.append(("float additive mask", (q, k, v), dict(attn_mask=float_mask), {})) checks.append(("smooth k + hadamard", (q, k, v), {}, dict(smooth_k=True, use_hadamard=True))) checks.append(("smooth + hadamard + int8 pv", (q, k, v), {}, dict(smooth_k=True, use_hadamard=True, pv_matmul_dtype="int8"))) qp, kp, vp = make_qkv(2, 8, 512, 40, structured=False) checks.append(("head dim 40 (padded)", (qp, kp, vp), {}, {})) checks.append(("head dim 40 + hadamard", (qp, kp, vp), {}, dict(use_hadamard=True))) qn, kn, vn = make_qkv(2, 8, 1000, 64, structured=False) checks.append(("non-pow2 sequence 1000", (qn, kn, vn), {}, {})) qx, kx, vx = make_qkv(2, 8, 512, 64, structured=False, kv_tokens=77) checks.append(("cross-attention 512q 77kv", (qx, kx, vx), {}, {})) # stress rows model the extreme-activation class (qwen-image, z-image base) where quantized # attention paths have produced non-finite outputs in the wild; they pass on finite output, # errors are informational since saturated softmax legitimately degrades accuracy qs, ks, vs = make_qkv(2, 8, 512, 64, structured=False) checks.append(("stress: 100x activations", (qs * 100.0, ks * 100.0, vs * 100.0), {}, {})) ko = ks.clone() ko[:, :, :2, :] *= 1000.0 checks.append(("stress: 1000x outlier keys", (qs, ko, vs), {}, {})) checks.append(("stress: fp16 100x activations", ((qs * 100.0).to(torch.float16), (ks * 100.0).to(torch.float16), (vs * 100.0).to(torch.float16)), {}, {})) # value-side stress: the existing rows perturb k (the qk-quant failure axis); the unsafe # pv accumulation (the accum flag with an unquantized pv, the sage-style fast mode) fails # on v instead. unnormalized softmax weights are <= 1 each until the epilogue normalize, # so near-uniform scores with same-sign values push one 16-wide kv-block dot past the # fp16 max: those rows pass by overflowing, while the scaled paths must stay finite vo = vs.clone() vo[:, :, :2, :] *= 1000.0 checks.append(("stress: 1000x outlier values", (qs, ks, vo), {}, {})) qu = qs * 1e-3 # near-zero scores: every unnormalized softmax weight ~1 vu = vs.abs() * 10000.0 # same-sign, below the fp16 cast limit checks.append(("stress: uniform attention 10000x values", (qu, ks, vu), {}, {})) if atten_supports_fp16_accum(): checks.append(("stress: uniform 10000x values, scaled fp16 accum", (qu, ks, vu), {}, dict(matmul_dtype="float16", pv_matmul_dtype="float16", use_fp16_accum=True))) checks.append(("stress: uniform 10000x values, unsafe pv accum", (qu, ks, vu), {}, dict(pv_matmul_dtype="disabled", use_fp16_accum=True))) checks.append(("stress: bf16 100000x values, unsafe pv accum", (qs.to(torch.bfloat16), ks.to(torch.bfloat16), vs.to(torch.bfloat16) * 100000.0), {}, dict(pv_matmul_dtype="disabled", use_fp16_accum=True))) table = Table(box=box.SIMPLE_HEAVY) table.add_column("code path") table.add_column("kernel error", justify="right") table.add_column("int8 error", justify="right") table.add_column("int8 max token", justify="right") table.add_column("result", justify="center") panel = Panel(table, title="kernel correctness", subtitle="[dim]small shapes, vs fp32 sdpa reference; stress rows pass on finite output, unsafe accum rows by overflowing[/dim]", box=ROUNDED_BOX, expand=False) failed = [] details = {} # checks target kernel behavior, so the input prep runs eager via a module-global swap. # do not toggle torch._dynamo.config.disable for this: newer torch raises "found no # compiled frames" when a fullgraph-compiled function is called inside a disable window, # failing every check and poisoning the first compiled call afterwards from sdnq.kernels import triton_atten as atten_module compiled_prep = atten_module.get_attn_inputs inner_prep = getattr(compiled_prep, "_torchdynamo_orig_callable", None) if inner_prep is not None: atten_module.get_attn_inputs = inner_prep aborted_after = None progress, task = live_progress() try: with Live(Group(panel, progress), console=console, refresh_per_second=4) as live: for index, (name, (cq, ck, cv), kwargs, sdnq_kwargs) in enumerate(checks, start=1): progress.reset(task, description=f"check {index}/{len(checks)} {name}") try: ref_kwargs = dict(kwargs) if ref_kwargs.get("attn_mask", None) is not None and torch.is_floating_point(ref_kwargs["attn_mask"]): ref_kwargs["attn_mask"] = ref_kwargs["attn_mask"].to(torch.float32) # stock sdpa requires the additive mask dtype to match query ref = fp32_reference(cq, ck, cv, **ref_kwargs) plain = sdnq_triton_atten(cq, ck, cv, do_quantize=False, **kwargs) quant_kwargs = dict(matmul_dtype="auto", pv_matmul_dtype="auto", **kwargs) quant_kwargs.update(sdnq_kwargs) quant = sdnq_triton_atten(cq, ck, cv, **quant_kwargs) err_plain = rel_err(plain, ref) err_quant = rel_err(quant, ref) max_err = max_token_err(quant, ref) finite_quant = bool(torch.isfinite(quant).all().item()) finite = bool(torch.isfinite(plain).all().item()) and finite_quant expect_overflow = name.endswith("unsafe pv accum") if expect_overflow: ok = not finite_quant # the row demonstrates the overflow; finite output means the demonstration failed, not that the mode is safe elif name.startswith("stress:"): ok = finite else: ok = finite and err_plain < 0.01 and err_quant < 0.2 if not ok: failed.append(name) details[name] = dict(kernel_err=err_plain, int8_err=err_quant, max_token_err=max_err, finite=finite, ok=ok) if expect_overflow: verdict = "[yellow]overflows (as designed)[/yellow]" if ok else "[red]did not overflow[/red]" elif not finite: verdict = "[red]non-finite[/red]" else: verdict = "[green]pass[/green]" if ok else "[red]fail[/red]" table.add_row(name, err_cell(err_plain), err_cell(err_quant), err_cell(max_err), verdict) except Exception as e: failed.append(name) details[name] = dict(error=error_summary(e, 200)) table.add_row(name, "-", "-", "-", failure_text(e)) if not cuda_context_alive(): aborted_after = name break live.update(panel) finally: atten_module.get_attn_inputs = compiled_prep console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(panel) if aborted_after is not None: emit(f"[red]the {aborted_after!r} check faulted the gpu and cuda errors are sticky: remaining checks and benchmarks cannot run in this process. rerun with --skip-checks and shapes that avoid the faulting geometry to collect the rest[/red]") report["correctness_aborted_after"] = aborted_after if failed: emit(f"[red]failed checks: {', '.join(failed)}[/red]") if "head dim 40 + hadamard" in failed and shared.opts.sdnq_attention_use_hadamard: emit("[yellow]current config has Use Hadamard enabled: models with non power-of-2 head dims (SD 1.5) fail at generation with it[/yellow]") report["correctness"] = dict(failed=failed, total=len(checks), checks=details) return failed def make_prep_fn(q, k, v, attn_mask, kwargs, is_causal=False, enable_gqa=False): # mirror sdnq_triton_atten's prep call so the prep column measures the same code path from sdnq.kernels import triton_atten as atten_module from sdnq.quant_utils import get_hadamard, get_hadamard_group_size from sdnq.utils import next_power_of_2 matmul_dtype = kwargs.get("matmul_dtype", "int8") do_quantize = kwargs.get("do_quantize", True) hadamard_group_size = kwargs.get("hadamard_group_size", 256) hadamard = None if kwargs.get("use_hadamard", False) and do_quantize and matmul_dtype not in {None, "none", "no"}: channel_size = next_power_of_2(min(q.shape[-1], k.shape[-1])) hadamard_group_size = min(hadamard_group_size, channel_size) enabled, hadamard_group_size = get_hadamard_group_size(channel_size, hadamard_group_size) if enabled: hadamard = get_hadamard(hadamard_group_size, dtype=q.dtype, device=q.device) def prep(): return atten_module.get_attn_inputs( # module lookup so the static-compile patch applies query=q, key=k, value=v, hadamard=hadamard, attn_mask=attn_mask, dropout_p=0.0, is_causal=is_causal, scale=None, enable_gqa=enable_gqa, smooth_k=kwargs.get("smooth_k", False), hadamard_group_size=hadamard_group_size, matmul_dtype=matmul_dtype, pv_matmul_dtype=kwargs.get("pv_matmul_dtype", None), do_quantize=do_quantize, out_dtype=None, ) return prep def bench_shape(preset, iters, warmup, position=None, config_timeout=None, fp8_result=None, selected=None): preset_cfg = shape_presets[preset] config_timeout = resolve_timeout(config_timeout, preset_cfg.get("config_timeout")) batch, heads, tokens, head_dim = preset_cfg["batch"], preset_cfg["heads"], preset_cfg["tokens"], preset_cfg["head_dim"] kv_tokens = preset_cfg.get("kv_tokens", tokens) kv_heads = preset_cfg.get("kv_heads", heads) causal = preset_cfg.get("causal", False) gqa = kv_heads != heads description = preset_cfg["desc"] iters = preset_cfg.get("iters", iters) warmup = preset_cfg.get("warmup", warmup) ref_head_chunk = preset_cfg.get("ref_head_chunk", 0) excluded_configs = preset_excluded_configs.get(preset, set()) if preset == "sd15": emit("[yellow]sd15: hadamard configs skipped, compiling hadamard with a non pow2 head dim currently hangs torch inductor[/yellow]") attn_mask = build_preset_mask(preset_cfg, torch_device) mask_nan_guard = preset_cfg.get("mask_nan_guard", False) sage = sage_attention() sage_fp16 = sage_attention_fp16_accum() amd_flash = amd_triton_flash() selected_configs = [] for config_id, label, kwargs in bench_configs: if config_id in excluded_configs or (selected is not None and config_id not in selected): continue if config_id == "sage" and (sage is None or attn_mask is not None or head_dim not in {64, 96, 128} or kv_tokens != tokens or gqa or causal): continue if config_id == "sagefp16" and (sage_fp16 is None or attn_mask is not None or head_dim not in {64, 96, 128} or kv_tokens != tokens or gqa or causal): continue if config_id == "amdflash" and (amd_flash is None or attn_mask is not None or head_dim > 128 or gqa): continue if config_id.startswith("flex") and (not flex_available() or attn_mask is not None or causal or kv_tokens != tokens): continue # a block only mask cannot carry a token mask or a causal rule, and cross attention is not sparsified if is_sdnq_sparse(config_id) and (not atten_supports_block_mask() or causal or kv_tokens != tokens): continue # the kernel composes a token mask with the block mask, so only the causal and cross attention rules apply if config_id == "fp8qk" and not (fp8_result and fp8_result["qk"][0]): continue if config_id == "fp8pv" and not (fp8_result and fp8_result["pv"][0]): continue if config_id == "fp8full" and not (fp8_result and fp8_result["qk"][0] and fp8_result["pv"][0]): continue if kwargs is not None and kwargs.get("use_fp16_accum") and not atten_supports_fp16_accum(): continue selected_configs.append((config_id, config_label(config_id, label), kwargs)) def make_table(): shape_table = Table(box=box.SIMPLE_HEAVY) shape_table.add_column("config") shape_table.add_column("median time", justify="right") shape_table.add_column("prep", justify="right") shape_table.add_column("speedup", justify="right") shape_table.add_column("error", justify="right") return shape_table geometry = f"batch={batch} heads={heads} tokens={tokens} head_dim={head_dim}" if kv_tokens != tokens: geometry += f" kv_tokens={kv_tokens}" if gqa: geometry += f" kv_heads={kv_heads}" if causal: geometry += " causal" def make_panel(content): return Panel(content, title=f"{preset}: {geometry} {dtype_label()}", subtitle=f"[dim]{description}[/dim]", box=ROUNDED_BOX, expand=False) table = make_table() panel = make_panel(table) results = {} rows = [] base_ms = None prefix = f"shape {position[0]}/{position[1]} " if position else "" progress, task = live_progress() with Live(Group(panel, progress), console=console, refresh_per_second=4) as live: progress.update(task, description=f"{prefix}{preset}: preparing inputs and fp32 reference") q, k, v = make_qkv(batch, heads, tokens, head_dim, kv_heads=kv_heads, kv_tokens=kv_tokens) scale = head_dim ** -0.5 ref = fp32_reference(q, k, v, attn_mask=attn_mask, is_causal=causal, enable_gqa=gqa, head_chunk=ref_head_chunk) if mask_nan_guard: ref = torch.nan_to_num(ref) anchor_fn = None for index, (config_id, label, kwargs) in enumerate(selected_configs, start=1): def phase(step, current_label=label, current_index=index): progress.update(task, description=f"{prefix}config {current_index}/{len(selected_configs)} {current_label}: {step}") if config_id == "base": def fn(mask=attn_mask): return torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=causal, enable_gqa=gqa) elif config_id == "sage": def fn(sm=scale): return sage(q, k, v, sm) elif config_id == "sagefp16": def fn(sm=scale): return sage_fp16(q, k, v, sm) elif config_id == "amdflash": def fn(sm=scale): return amd_flash(q, k, v, sm, is_causal=causal) elif config_id.startswith("flex"): fn = make_flex_fn(config_id, q, k, v, scale, gqa) elif is_sdnq_sparse(config_id): fn = make_sdnq_sparse_fn(config_id, q, k, v, attn_mask, kwargs, causal, gqa) else: def fn(kw=kwargs, mask=attn_mask): return sdnq_triton_atten(q, k, v, attn_mask=mask, is_causal=causal, enable_gqa=gqa, **kw) try: torch._dynamo.reset() # pylint: disable=protected-access # one compiled input-prep specialization per config, matching how the webui runs a fixed config progress.reset(task) # restart the elapsed clock so it times the current config phase("compiling") with time_limit(config_timeout, label): err = rel_err(torch.nan_to_num(fn()) if mask_nan_guard else fn(), ref) ms, ms_sigma = bench_stats(fn, warmup, iters, on_phase=phase) prep_cell = "-" prep_ms = None if kwargs is not None: # time the input prep alone; included in the median try: phase("timing input prep") prep_ms = bench(make_prep_fn(q, k, v, attn_mask, kwargs, is_causal=causal, enable_gqa=gqa), warmup, iters) prep_cell = f"[dim]{prep_ms:8.3f} ms[/dim]" except Exception: pass if base_ms is None: base_ms = ms if config_id == "int8": anchor_fn = fn results[config_id] = dict(ms=ms, ms_sigma=ms_sigma, err=err, prep_ms=prep_ms) row = (label, f"{ms:8.3f} ms", prep_cell, speedup_cell(base_ms, ms), f"{err:.5f}") rows.append((config_id, ms, row)) table.add_row(*row) except Exception as e: results[config_id] = dict(ms=None, err=None, prep_ms=None, error=error_summary(e, 200)) row = (label, "-", "-", "-", failure_text(e)) rows.append((config_id, None, row)) table.add_row(*row) # sentinel: re-time the anchor config minutes after its first measurement; the # delta samples run-level clock drift, which feeds the verdict noise floor anchor_first = (results.get("int8") or {}).get("ms") if anchor_fn is not None and anchor_first: try: progress.update(task, description=f"{prefix}{preset}: drift sentinel, re-timing int8") record_drift(anchor_first, bench(anchor_fn, warmup, iters)) except Exception: pass # rebuild the table to star the best sdnq config when it beats the baseline best_id = best_config(results) if base_ms and best_id is not None and results[best_id]["ms"] < base_ms: table = make_table() for config_id, _ms, row in rows: if config_id == best_id: table.add_row(f"* {row[0]}", *row[1:], style="bold") # ascii marker: legacy windows consoles crash rendering non-cp1252 characters else: table.add_row(*row) panel = make_panel(table) live.update(panel) console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(panel) if dtype_label() == "bfloat16": for (band_shape, band_config), (band_low, band_high) in error_sanity_bands.items(): band_err = (results.get(band_config) or {}).get("err") if band_shape == preset else None if band_err and not band_low <= band_err <= band_high: emit(f"[yellow]measurement sanity: {preset} {band_config} error {band_err:.5f} sits outside the cross-gpu band [{band_low:.4f}, {band_high:.4f}]; treat this run's numbers with suspicion[/yellow]") report.setdefault("attention", {})[preset] = dict(geometry=geometry, results=results) return results def bench_dequant_shape(shape_label, out_features, in_features, iters, warmup, position=None, config_timeout=300, selected_dtypes=None): from sdnq.common import check_torch_compile compile_on = check_torch_compile() dtype_configs = [(dtype_id, label, cfg) for dtype_id, label, cfg in dequant_dtype_configs if selected_dtypes is None or dtype_id in selected_dtypes] def make_table(): shape_table = Table(box=box.SIMPLE_HEAVY) shape_table.add_column("config") shape_table.add_column("quant", justify="right") shape_table.add_column("size", justify="right") shape_table.add_column("deq eager", justify="right") shape_table.add_column("deq compiled", justify="right") shape_table.add_column("weight err", justify="right") shape_table.add_column("fwd eager", justify="right") shape_table.add_column("fwd compiled", justify="right") shape_table.add_column("quantized mm", justify="right") shape_table.add_column("mm err", justify="right") shape_table.add_column("speedup", justify="right") return shape_table def make_panel(content): subtitle = f"[dim]{dequant_forward_tokens} token input, bias-free linears, errors vs fp32; quantized mm and speedup follow the current config (compile {'on' if compile_on else 'off'})[/dim]" return Panel(content, title=f"weight dequant: {shape_label} {dtype_label()}", subtitle=subtitle, box=ROUNDED_BOX, expand=False) table = make_table() panel = make_panel(table) results = {} notes = [] prefix = f"shape {position[0]}/{position[1]} " if position else "" progress, task = live_progress() with Live(Group(panel, progress), console=console, refresh_per_second=4) as live: progress.update(task, description=f"{prefix}{shape_label}: preparing weights and bf16 baseline") weight = make_source_weight(out_features, in_features) weight_fp32 = weight.to(torch.float32) generator = torch.Generator(device=torch_device).manual_seed(42) x = torch.randn(dequant_forward_tokens, in_features, device=torch_device, dtype=bench_dtype, generator=generator) baseline = torch.nn.Linear(in_features, out_features, bias=False, device=torch_device, dtype=bench_dtype) with torch.no_grad(): baseline.weight.copy_(weight) def baseline_fn(): return baseline(x) ref_out = fp32_linear_reference(x, weight_fp32) base_ms, base_ms_sigma = bench_stats(baseline_fn, warmup, iters) base_err = rel_err(baseline_fn(), ref_out) base_bytes = baseline.weight.numel() * baseline.weight.element_size() results["bf16"] = dict(fwd_ms=base_ms, fwd_ms_sigma=base_ms_sigma, fwd_err=base_err, size_bytes=base_bytes) table.add_row(f"{dtype_label()} nn.Linear", "-", size_cell(base_bytes), "-", "-", "-", f"{base_ms:8.3f} ms", "-", "-", "-", "x1.00") for index, (dtype_id, label, cfg) in enumerate(dtype_configs, start=1): def phase(step, current_label=label, current_index=index): progress.update(task, description=f"{prefix}config {current_index}/{len(dtype_configs)} {current_label}: {step}") entry = dict(eager_ms=None, compiled_ms=None, fwd_eager_ms=None, fwd_compiled_ms=None, fwd_err=None, mm_ms=None, mm_err=None, mm_dtype=None, weight_err=None, quant_s=None, size_bytes=None) results[dtype_id] = entry progress.reset(task) phase("quantizing") try: layer, entry["quant_s"] = make_quantized_linear(weight, **cfg) entry["size_bytes"] = layer_storage_bytes(layer) except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(label, "-", "-", "-", "-", "-", "-", "-", "-", "-", failure_text(e)) continue def eager_fn(current=layer): return current.sdnq_dequantizer( current.weight, current.scale, zero_point=getattr(current, "zero_point", None), svd_up=getattr(current, "svd_up", None), svd_down=getattr(current, "svd_down", None), skip_compile=True, ) try: phase("timing eager dequant") eager_out = eager_fn().to(torch.float32) entry["eager_ms"], entry["eager_ms_sigma"] = bench_stats(eager_fn, warmup, iters) entry["weight_err"] = rel_err(eager_out, weight_fp32) eager_cell = f"{entry['eager_ms']:8.3f} ms" except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(label, quant_cell(entry["quant_s"]), size_cell(entry["size_bytes"]), failure_text(e), "-", "-", "-", "-", "-", "-", "-") del layer continue try: args, kwargs = dequant_args(layer) def compiled_fn(current_args=args, current_kwargs=kwargs): return get_compiled_dequantize_weight()(*current_args, **current_kwargs) phase("compiling dequant") with time_limit(config_timeout, label): compiled_out = compiled_fn() torch_device_module.synchronize() entry["compiled_ms"], entry["compiled_ms_sigma"] = bench_stats(compiled_fn, warmup, iters, on_phase=phase) compiled_cell = f"{entry['compiled_ms']:8.3f} ms" drift = rel_err(compiled_out, eager_out) if drift > 1e-3: notes.append(f"[yellow]{label}: compiled vs eager dequant drift {drift:.2e}[/yellow]") entry["compiled_drift"] = drift del compiled_out except Exception as e: entry["compiled_error"] = error_summary(e, 200) compiled_cell = failure_text(e) reset_compiled_dequant() del eager_out def fwd_fn(current=layer): return current(x) # eager-mode forward: the dequantizer's __call__ resolves dequantize_weight_compiled # as a module global at call time, so pointing it at the eager function for the # bench matches the webui with sdnq torch.compile off (SDNQ_USE_TORCH_COMPILE=0). do not toggle # torch._dynamo.config.disable instead: code objects called during a disable window # keep their skip marking and never compile again in this process from sdnq import dequantizer as dequantizer_module saved_compiled_fn = dequantizer_module.dequantize_weight_compiled try: phase("timing linear forward, eager dequant") dequantizer_module.dequantize_weight_compiled = dequantizer_module.dequantize_weight fwd_fn() torch_device_module.synchronize() entry["fwd_eager_ms"], entry["fwd_eager_ms_sigma"] = bench_stats(fwd_fn, warmup, iters) entry["fwd_err"] = rel_err(fwd_fn(), ref_out) fwd_eager_cell = f"{entry['fwd_eager_ms']:8.3f} ms" except Exception as e: entry["fwd_eager_error"] = error_summary(e, 200) fwd_eager_cell = failure_text(e) finally: dequantizer_module.dequantize_weight_compiled = saved_compiled_fn try: phase("compiling linear forward") with time_limit(config_timeout, label): fwd_fn() torch_device_module.synchronize() phase("timing linear forward, compiled dequant") entry["fwd_compiled_ms"], entry["fwd_compiled_ms_sigma"] = bench_stats(fwd_fn, warmup, iters) fwd_compiled_cell = f"{entry['fwd_compiled_ms']:8.3f} ms" except Exception as e: entry["fwd_compiled_error"] = error_summary(e, 200) fwd_compiled_cell = failure_text(e) torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state del layer fwd_ms = entry["fwd_compiled_ms"] if compile_on else entry["fwd_eager_ms"] entry["fwd_ms"] = fwd_ms # the mode the current config runs; recommendations and speedup key off it entry["fwd_ms_sigma"] = entry.get("fwd_compiled_ms_sigma") if compile_on else entry.get("fwd_eager_ms_sigma") speedup = speedup_cell(base_ms, fwd_ms) if fwd_ms else "-" try: phase("quantizing for quantized matmul") mm_layer, _mm_quant_seconds = make_quantized_linear(weight, use_quantized_matmul=True, **cfg) entry["mm_dtype"] = mm_layer.sdnq_dequantizer.quantized_matmul_dtype def mm_fn(current=mm_layer): return current(x) phase("timing quantized matmul forward") with time_limit(config_timeout, label): mm_fn() torch_device_module.synchronize() entry["mm_ms"], entry["mm_ms_sigma"] = bench_stats(mm_fn, warmup, iters) entry["mm_err"] = rel_err(mm_fn(), ref_out) mm_cell = f"{entry['mm_ms']:8.3f} ms [dim]{entry['mm_dtype']}[/dim]" del mm_layer except Exception as e: entry["mm_error"] = error_summary(e, 200) mm_cell = failure_text(e) if entry["fwd_err"] and entry["weight_err"] and abs(entry["fwd_err"] - entry["weight_err"]) > 0.2 * entry["weight_err"]: notes.append(f"[yellow]{label}: fwd err {entry['fwd_err']:.5f} diverges from weight err {entry['weight_err']:.5f}[/yellow]") table.add_row(label, quant_cell(entry["quant_s"]), size_cell(entry["size_bytes"]), eager_cell, compiled_cell, err_cell(entry["weight_err"]), fwd_eager_cell, fwd_compiled_cell, mm_cell, err_cell(entry["mm_err"]), speedup) # sentinel: re-time the bf16 baseline to sample run-level clock drift for this shape window try: progress.update(task, description=f"{prefix}{shape_label}: drift sentinel, re-timing the {dtype_label()} baseline") record_drift(base_ms, bench(baseline_fn, warmup, iters)) except Exception: pass live.update(panel) console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(panel) for note in notes: emit(note) report.setdefault("dequant", {})[shape_label] = dict(out_features=out_features, in_features=in_features, base_ms=base_ms, results=results) return results def bench_dequant_variants(shape_label, out_features, in_features, plain_results, selected_dtypes, selected_variants, iters, warmup, config_timeout=300): # svd and hadamard on top of the base dtypes: the quantize-time, size and error cost of # each option, benched through the production dequantizer path at one layer shape variant_configs = [(variant_id, cfg) for variant_id, cfg in dequant_variant_configs if variant_id in selected_variants] dtype_configs = [(dtype_id, label, cfg) for dtype_id, label, cfg in dequant_dtype_configs if dtype_id in dequant_variant_dtypes and (selected_dtypes is None or dtype_id in selected_dtypes)] if not variant_configs or not dtype_configs: return {} table = Table(box=box.SIMPLE_HEAVY) table.add_column("config") table.add_column("quant", justify="right") table.add_column("size", justify="right") table.add_column("deq compiled", justify="right") table.add_column("weight err", justify="right") table.add_column("err vs plain", justify="right") table.add_column("fwd compiled", justify="right") table.add_column("speedup", justify="right") panel = Panel(table, title=f"svd / hadamard variants: {shape_label} {dtype_label()}", subtitle="[dim]plain rows repeated dimmed for comparison; err vs plain below x1.00 = better reconstruction[/dim]", box=ROUNDED_BOX, expand=False) results = {} base_ms = plain_results.get("bf16", {}).get("fwd_ms") runs = [(dtype_id, label, cfg, variant_id, variant_cfg) for dtype_id, label, cfg in dtype_configs for variant_id, variant_cfg in variant_configs] progress, task = live_progress() with Live(Group(panel, progress), console=console, refresh_per_second=4) as live: weight = make_source_weight(out_features, in_features) weight_fp32 = weight.to(torch.float32) generator = torch.Generator(device=torch_device).manual_seed(42) x = torch.randn(dequant_forward_tokens, in_features, device=torch_device, dtype=bench_dtype, generator=generator) added_plain = set() for index, (dtype_id, label, cfg, variant_id, variant_cfg) in enumerate(runs, start=1): row_label = f"{label} + {variant_id}" def phase(step, current_label=row_label, current_index=index): progress.update(task, description=f"variant {current_index}/{len(runs)} {current_label}: {step}") if dtype_id not in added_plain: added_plain.add(dtype_id) plain = plain_results.get(dtype_id) or {} if plain.get("weight_err") is not None: plain_fwd = plain.get("fwd_compiled_ms") table.add_row( f"[dim]{label}[/dim]", f"[dim]{quant_cell(plain.get('quant_s'))}[/dim]", f"[dim]{size_cell(plain['size_bytes'])}[/dim]", f"[dim]{plain['compiled_ms']:8.3f} ms[/dim]" if plain.get("compiled_ms") else "-", f"[dim]{plain['weight_err']:.5f}[/dim]", "[dim]x1.00[/dim]", f"[dim]{plain_fwd:8.3f} ms[/dim]" if plain_fwd else "-", f"[dim]{speedup_cell(base_ms, plain_fwd)}[/dim]" if base_ms and plain_fwd else "-", ) entry = dict(quant_s=None, size_bytes=None, compiled_ms=None, weight_err=None, err_ratio=None, fwd_ms=None) results[f"{dtype_id}+{variant_id}"] = entry progress.reset(task) phase("quantizing") try: layer, entry["quant_s"] = make_quantized_linear(weight, **cfg, **variant_cfg) entry["size_bytes"] = layer_storage_bytes(layer) except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(row_label, "-", "-", "-", "-", "-", "-", failure_text(e)) continue def deq_fn(current=layer): return current.sdnq_dequantizer( current.weight, current.scale, zero_point=getattr(current, "zero_point", None), svd_up=getattr(current, "svd_up", None), svd_down=getattr(current, "svd_down", None), ) def fwd_fn(current=layer): return current(x) try: phase("compiling dequant") with time_limit(config_timeout, row_label): deq_out = deq_fn().to(torch.float32) torch_device_module.synchronize() entry["compiled_ms"] = bench(deq_fn, warmup, iters, on_phase=phase) entry["weight_err"] = rel_err(deq_out, weight_fp32) plain_err = (plain_results.get(dtype_id) or {}).get("weight_err") entry["err_ratio"] = entry["weight_err"] / plain_err if plain_err else None del deq_out phase("timing linear forward") with time_limit(config_timeout, row_label): fwd_fn() torch_device_module.synchronize() entry["fwd_ms"] = bench(fwd_fn, warmup, iters) table.add_row( row_label, quant_cell(entry["quant_s"]), size_cell(entry["size_bytes"]), f"{entry['compiled_ms']:8.3f} ms", err_cell(entry["weight_err"]), f"x{entry['err_ratio']:.2f}" if entry["err_ratio"] else "-", f"{entry['fwd_ms']:8.3f} ms", speedup_cell(base_ms, entry["fwd_ms"]) if base_ms else "-", ) except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(row_label, quant_cell(entry["quant_s"]), size_cell(entry["size_bytes"]), "-", "-", "-", "-", failure_text(e)) torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state del layer live.update(panel) console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(panel) torch_device_module.empty_cache() report["dequant_variants"] = dict(shape=shape_label, results=results) return results def bench_float_mm_alternatives(shape_label, out_features, in_features, plain_results, selected_dtypes, iters, warmup, config_timeout=300): # explicit MatMul type for float storage dtypes: under enabled these route quantized matmul # to fp8, which not every gpu can run; measure what setting int8 or float16 costs instead dtype_configs = [(dtype_id, label, cfg) for dtype_id, label, cfg in dequant_dtype_configs if dtype_id in float_mm_dtypes and (selected_dtypes is None or dtype_id in selected_dtypes)] if not dtype_configs: return {} table = Table(box=box.SIMPLE_HEAVY) table.add_column("config") table.add_column("quant", justify="right") table.add_column("mm fwd", justify="right") table.add_column("out err", justify="right") table.add_column("vs dequant", justify="right") panel = Panel(table, title=f"float weights, explicit MatMul type: {shape_label} {dtype_label()}", subtitle="[dim]quantized matmul with the MatMul type set explicitly; dequant path and enabled rows repeated dimmed; vs dequant above x1.00 = faster than the dequant-path forward[/dim]", box=ROUNDED_BOX, expand=False) results = {} # the accum variant repeats the float16 mm with the triton fp16-accum globals flipped # in-process; the env var is not a dropdown value, so its rows never enter candidacy mm_variants = [(mm_dtype, False) for mm_dtype in float_mm_alternative_dtypes] if triton_mm_supports_fp16_accum() and "float16" in float_mm_alternative_dtypes: mm_variants.append(("float16", True)) runs = [(dtype_id, label, cfg, mm_dtype, accum) for dtype_id, label, cfg in dtype_configs for mm_dtype, accum in mm_variants] progress, task = live_progress() with Live(Group(panel, progress), console=console, refresh_per_second=4) as live: weight = make_source_weight(out_features, in_features) weight_fp32 = weight.to(torch.float32) generator = torch.Generator(device=torch_device).manual_seed(42) x = torch.randn(dequant_forward_tokens, in_features, device=torch_device, dtype=bench_dtype, generator=generator) ref_out = fp32_linear_reference(x, weight_fp32) added_plain = set() for index, (dtype_id, label, cfg, mm_dtype, accum) in enumerate(runs, start=1): row_label = f"{label} + {mm_dtype} mm" + (", fp16 accum" if accum else "") def phase(step, current_label=row_label, current_index=index): progress.update(task, description=f"float mm {current_index}/{len(runs)} {current_label}: {step}") plain = plain_results.get(dtype_id) or {} fwd_ms = plain.get("fwd_ms") if dtype_id not in added_plain: added_plain.add(dtype_id) if fwd_ms: table.add_row(f"[dim]{label} dequant path[/dim]", "-", f"[dim]{fwd_ms:8.3f} ms[/dim]", f"[dim]{plain['fwd_err']:.5f}[/dim]" if plain.get("fwd_err") else "-", "[dim]x1.00[/dim]") if plain.get("mm_ms"): table.add_row(f"[dim]{label} + enabled ({plain.get('mm_dtype')})[/dim]", "-", f"[dim]{plain['mm_ms']:8.3f} ms[/dim]", f"[dim]{plain['mm_err']:.5f}[/dim]" if plain.get("mm_err") else "-", f"[dim]{speedup_cell(fwd_ms, plain['mm_ms'])}[/dim]" if fwd_ms else "-") elif plain.get("mm_error"): table.add_row(f"[dim]{label} + enabled ({plain.get('mm_dtype') or 'float8_e4m3fn'})[/dim]", "-", failure_text(RuntimeError(plain["mm_error"])), "-", "-") entry = dict(quant_s=None, mm_ms=None, mm_err=None, mm_dtype=None, fp16_accum=accum) results[f"{dtype_id}+{mm_dtype}" + ("+accum" if accum else "")] = entry progress.reset(task) phase("quantizing") try: layer, entry["quant_s"] = make_quantized_linear(weight, use_quantized_matmul=True, quantized_matmul_dtype=mm_dtype, **cfg) entry["mm_dtype"] = layer.sdnq_dequantizer.quantized_matmul_dtype except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(row_label, "-", failure_text(e), "-", "-") continue def mm_fn(current=layer, use_accum=accum): if use_accum: with triton_mm_fp16_accum(): return current(x) return current(x) try: phase("compiling quantized matmul forward") with time_limit(config_timeout, row_label): mm_fn() torch_device_module.synchronize() phase("timing quantized matmul forward") entry["mm_ms"], entry["mm_ms_sigma"] = bench_stats(mm_fn, warmup, iters) entry["mm_err"] = rel_err(mm_fn(), ref_out) table.add_row(row_label, quant_cell(entry["quant_s"]), f"{entry['mm_ms']:8.3f} ms", err_cell(entry["mm_err"]), speedup_cell(fwd_ms, entry["mm_ms"]) if fwd_ms else "-") except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(row_label, quant_cell(entry["quant_s"]), failure_text(e), "-", "-") torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state del layer live.update(panel) console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(panel) torch_device_module.empty_cache() report["dequant_float_mm"] = dict(shape=shape_label, results=results) return results # === quantized matmul backends === # # The linear layers bind their scaled-mm function at import # (`from ...kernel_wrappers import int_scaled_mm_func`), so SDNQ_USE_TRITON_MM freezes the # backend for the process and a cross-process A/B carries clock drift into the comparison. # The call sites look the name up as a module global at call time, so rebinding it on the # consuming module switches backends in-process; the layer forwards are compile_func'd, so # every swap needs a dynamo reset or the traced graph keeps calling the previous function. # # The torch row is whatever kernel_wrappers bound when triton was not selected, captured # rather than reimplemented. Run without SDNQ_USE_TRITON_MM=1 to have it available: where # triton is the platform default (xpu, ipex, zluda, rdna2 and older) the torch fallbacks are # never defined and the row needs SDNQ_USE_TRITON_MM=0. mm_swap_targets = [ ("sdnq.layers.linear.linear_int8", "int_scaled_mm_func"), ("sdnq.layers.linear.linear_uint8", "int_scaled_mm_func"), ("sdnq.layers.linear.linear_fp16", "fp_scaled_mm_func"), ("sdnq.layers.linear.linear_fp8", "fp8_scaled_mm_func"), ] def mm_backend_bindings(): # {backend: {(module, attr): func}} for the backends bindable in this process, plus a # note for any that are not bound = {} for module_path, attr in mm_swap_targets: try: module = importlib.import_module(module_path) except Exception: continue func = getattr(module, attr, None) if func is not None: bound[(module_path, attr)] = func try: from sdnq.kernels.triton_scaled_mm import sdnq_scaled_mm except Exception as e: return {}, {"triton": f"triton scaled mm unavailable: {error_summary(e, 120)}"} backends, unavailable = {}, {} if bound and all(func is sdnq_scaled_mm for func in bound.values()): unavailable["torch"] = "triton is the default matmul backend on this device; rerun with SDNQ_USE_TRITON_MM=0 to bind the torch fallbacks" elif bound: backends["torch"] = dict(bound) backends["triton"] = {target: sdnq_scaled_mm for target in bound} return backends, unavailable def apply_mm_backend(binding): for (module_path, attr), func in binding.items(): setattr(importlib.import_module(module_path), attr, func) torch._dynamo.reset() # pylint: disable=protected-access # layer forwards are compiled: the traced graph pins the previous function def bench_mm_backends(shape_label, out_features, in_features, selected_dtypes, backends, iters, warmup, config_timeout=300, rounds=2): # paired same-run comparison of the quantized-matmul backends: one quantized layer per # dtype, benched through each backend in turn so both rows see the same weights and the # same clock state. Round order alternates so monotonic drift cancels instead of # accumulating into whichever backend runs second; each row keeps its fastest round. available, unavailable = mm_backend_bindings() selected = [name for name in backends if name in available] for name in backends: if name in unavailable: emit(f"[yellow]matmul backend '{name}' not benchable: {unavailable[name]}[/yellow]") if len(selected) < 2: if selected: emit(f"[yellow]matmul backend comparison needs two bindable backends, only '{selected[0]}' is available; skipping[/yellow]") return {} dtype_configs = [(dtype_id, label, cfg) for dtype_id, label, cfg in dequant_dtype_configs if selected_dtypes is None or dtype_id in selected_dtypes] if not dtype_configs: return {} # the comparison only runs when the torch row is bindable, which means the process came up on # it; restore it after the sweep so later sections measure the config the user actually runs original_binding = available["torch"] table = Table(box=box.SIMPLE_HEAVY) table.add_column("weights") for name in selected: table.add_column(f"{name} mm", justify="right") table.add_column("delta", justify="right") table.add_column("out err", justify="right") panel = Panel( table, title=f"quantized matmul backends, paired: {shape_label} {dtype_label()}", subtitle=f"[dim]same layer and clock state, {rounds} alternating rounds, fastest kept; delta = {selected[-1]} vs {selected[0]}, negative = {selected[-1]} faster[/dim]", box=ROUNDED_BOX, expand=False, ) results = {} progress, task = live_progress() with Live(Group(panel, progress), console=console, refresh_per_second=4) as live: weight = make_source_weight(out_features, in_features) weight_fp32 = weight.to(torch.float32) generator = torch.Generator(device=torch_device).manual_seed(42) x = torch.randn(dequant_forward_tokens, in_features, device=torch_device, dtype=bench_dtype, generator=generator) ref_out = fp32_linear_reference(x, weight_fp32) for index, (dtype_id, label, cfg) in enumerate(dtype_configs, start=1): def phase(step, current_label=label, current_index=index): progress.update(task, description=f"mm backends {current_index}/{len(dtype_configs)} {current_label}: {step}") entry = dict(mm_dtype=None, backends={}) results[dtype_id] = entry try: phase("quantizing for quantized matmul") layer, _quant_seconds = make_quantized_linear(weight, use_quantized_matmul=True, **cfg) entry["mm_dtype"] = layer.sdnq_dequantizer.quantized_matmul_dtype except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(label, *["-"] * len(selected), "-", failure_text(e)) continue def mm_fn(current=layer): return current(x) for round_index in range(rounds): order = selected if round_index % 2 == 0 else list(reversed(selected)) for name in order: slot = entry["backends"].setdefault(name, dict(ms=None, err=None)) if slot.get("error"): continue try: phase(f"{name} backend, round {round_index + 1}/{rounds}") apply_mm_backend(available[name]) with time_limit(config_timeout, f"{label} {name} mm"): mm_fn() torch_device_module.synchronize() ms = bench(mm_fn, warmup, iters) if slot["ms"] is None or ms < slot["ms"]: slot["ms"] = ms if slot["err"] is None: slot["err"] = rel_err(mm_fn(), ref_out) except Exception as e: slot["error"] = error_summary(e, 200) torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state del layer cells = [] for name in selected: slot = entry["backends"].get(name, {}) cells.append(f"{slot['ms']:8.3f} ms" if slot.get("ms") else failure_text(RuntimeError(slot.get("error", "not run")))) first, last = entry["backends"].get(selected[0], {}), entry["backends"].get(selected[-1], {}) if first.get("ms") and last.get("ms"): delta = (last["ms"] - first["ms"]) / first["ms"] * 100 colour = "green" if delta < -3 else ("red" if delta > 3 else "dim") entry["delta_pct"] = delta delta_cell = f"[{colour}]{delta:+.1f}%[/{colour}]" else: delta_cell = "-" errs = {slot.get("err") for slot in entry["backends"].values() if slot.get("err") is not None} err_text = err_cell(max(errs)) if errs else "-" if len(errs) > 1 and max(errs) - min(errs) > 1e-4: err_text += " [yellow]differs[/yellow]" # backends must be numerically equivalent; a split here is a kernel bug table.add_row(label, *cells, delta_cell, err_text) live.update(panel) console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(panel) apply_mm_backend(original_binding) torch_device_module.empty_cache() report["dequant_mm_backends"] = dict(shape=shape_label, backends=selected, rounds=rounds, results=results) return results def resolved_group_label(layer, in_features): # infer the group size sdnq actually used from the stored scale shape: grouped scales are # [out, groups, 1], row-wise scales collapse the group axis scale = getattr(layer, "scale", None) if scale is None: return "-" if scale.ndim >= 3 and scale.shape[1] > 1: return f"g{in_features // scale.shape[1]}" return "row" def bench_group_sizes(shape_label, out_features, in_features, plain_results, selected_dtypes, iters, warmup, config_timeout=300): # pylint: disable=unused-argument # the Group size setting: 0 = auto, -1 = row-wise, explicit values snap to a divisor of # in_features; grouping forces a per-forward re-quantize when quantized matmul is on, so # the mm cells price that cost alongside the accuracy gain dtype_configs = [(dtype_id, label, cfg) for dtype_id, label, cfg in dequant_dtype_configs if dtype_id in group_sweep_dtypes and (selected_dtypes is None or dtype_id in selected_dtypes)] if not dtype_configs: return {} table = Table(box=box.SIMPLE_HEAVY) table.add_column("config") table.add_column("quant", justify="right") table.add_column("size", justify="right") table.add_column("weight err", justify="right") table.add_column("fwd", justify="right") table.add_column("quantized mm", justify="right") table.add_column("mm err", justify="right") panel = Panel(table, title=f"group size sweep: {shape_label} {dtype_label()}", subtitle="[dim]Group size setting; auto and the mm path can resolve to different groups, resolved size shown per cell[/dim]", box=ROUNDED_BOX, expand=False) results = {} runs = [(dtype_id, label, cfg, group) for dtype_id, label, cfg in dtype_configs for group in group_sweep_values] progress, task = live_progress() with Live(Group(panel, progress), console=console, refresh_per_second=4) as live: weight = make_source_weight(out_features, in_features) weight_fp32 = weight.to(torch.float32) generator = torch.Generator(device=torch_device).manual_seed(42) x = torch.randn(dequant_forward_tokens, in_features, device=torch_device, dtype=bench_dtype, generator=generator) ref_out = fp32_linear_reference(x, weight_fp32) for index, (dtype_id, label, cfg, group) in enumerate(runs, start=1): group_label = {0: "auto", -1: "row"}.get(group, str(group)) row_label = f"{label} group {group_label}" def phase(step, current_label=row_label, current_index=index): progress.update(task, description=f"group {current_index}/{len(runs)} {current_label}: {step}") entry = dict(quant_s=None, size_bytes=None, weight_err=None, fwd_ms=None, mm_ms=None, mm_err=None, group=group, resolved=None, mm_resolved=None) results[f"{dtype_id}@{group}"] = entry progress.reset(task) phase("quantizing") base_cfg = {key: value for key, value in cfg.items() if key != "group_size"} try: layer, entry["quant_s"] = make_quantized_linear(weight, group_size=group, **base_cfg) entry["size_bytes"] = layer_storage_bytes(layer) entry["resolved"] = resolved_group_label(layer, in_features) except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(row_label, "-", "-", failure_text(e), "-", "-", "-") continue def fwd_fn(current=layer): return current(x) try: phase("timing dequant and forward") with time_limit(config_timeout, row_label): deq_out = layer.sdnq_dequantizer( layer.weight, layer.scale, zero_point=getattr(layer, "zero_point", None), svd_up=getattr(layer, "svd_up", None), svd_down=getattr(layer, "svd_down", None), ).to(torch.float32) fwd_fn() torch_device_module.synchronize() entry["weight_err"] = rel_err(deq_out, weight_fp32) del deq_out entry["fwd_ms"] = bench(fwd_fn, warmup, iters) except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(row_label, quant_cell(entry["quant_s"]), size_cell(entry["size_bytes"]), "-", failure_text(e), "-", "-") torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state del layer continue del layer mm_cell, mm_err_cell = "-", "-" try: phase("quantizing for quantized matmul") mm_layer, _mm_quant_seconds = make_quantized_linear(weight, group_size=group, use_quantized_matmul=True, **base_cfg) entry["mm_resolved"] = resolved_group_label(mm_layer, in_features) def mm_fn(current=mm_layer): return current(x) phase("timing quantized matmul forward") with time_limit(config_timeout, row_label): mm_fn() torch_device_module.synchronize() entry["mm_ms"], entry["mm_ms_sigma"] = bench_stats(mm_fn, warmup, iters) entry["mm_err"] = rel_err(mm_fn(), ref_out) mm_suffix = f" [dim]{entry['mm_resolved']}[/dim]" if entry["mm_resolved"] != entry["resolved"] else "" mm_cell = f"{entry['mm_ms']:8.3f} ms{mm_suffix}" mm_err_cell = err_cell(entry["mm_err"]) del mm_layer except Exception as e: entry["mm_error"] = error_summary(e, 200) mm_cell = failure_text(e) torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state shown_label = f"{row_label} [dim]({entry['resolved']})[/dim]" if group == 0 else row_label table.add_row(shown_label, quant_cell(entry["quant_s"]), size_cell(entry["size_bytes"]), err_cell(entry["weight_err"]), f"{entry['fwd_ms']:8.3f} ms", mm_cell, mm_err_cell) live.update(panel) console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(panel) torch_device_module.empty_cache() report["dequant_group_sizes"] = dict(shape=shape_label, results=results) return results def bench_svd_ranks(shape_label, out_features, in_features, plain_results, selected_dtypes, iters, warmup, config_timeout=300): # the SVD rank size setting: rank drives both outlier absorption and the rank x (in + out) # fp16 size overhead; plain no-svd rows repeated dimmed for comparison dtype_configs = [(dtype_id, label, cfg) for dtype_id, label, cfg in dequant_dtype_configs if dtype_id in svd_rank_sweep_dtypes and (selected_dtypes is None or dtype_id in selected_dtypes)] if not dtype_configs: return {} table = Table(box=box.SIMPLE_HEAVY) table.add_column("config") table.add_column("quant", justify="right") table.add_column("size", justify="right") table.add_column("weight err", justify="right") table.add_column("err vs plain", justify="right") table.add_column("fwd", justify="right") panel = Panel(table, title=f"svd rank sweep: {shape_label} {dtype_label()}", subtitle="[dim]SVD rank size setting at svd steps 8; err vs plain below x1.00 = better reconstruction[/dim]", box=ROUNDED_BOX, expand=False) results = {} runs = [(dtype_id, label, cfg, rank) for dtype_id, label, cfg in dtype_configs for rank in svd_rank_sweep_values] progress, task = live_progress() with Live(Group(panel, progress), console=console, refresh_per_second=4) as live: weight = make_source_weight(out_features, in_features) weight_fp32 = weight.to(torch.float32) generator = torch.Generator(device=torch_device).manual_seed(42) x = torch.randn(dequant_forward_tokens, in_features, device=torch_device, dtype=bench_dtype, generator=generator) added_plain = set() for index, (dtype_id, label, cfg, rank) in enumerate(runs, start=1): row_label = f"{label} + svd rank {rank}" def phase(step, current_label=row_label, current_index=index): progress.update(task, description=f"svd {current_index}/{len(runs)} {current_label}: {step}") plain = plain_results.get(dtype_id) or {} if dtype_id not in added_plain: added_plain.add(dtype_id) if plain.get("weight_err") is not None: plain_fwd = plain.get("fwd_ms") table.add_row(f"[dim]{label}[/dim]", f"[dim]{quant_cell(plain.get('quant_s'))}[/dim]", f"[dim]{size_cell(plain['size_bytes'])}[/dim]", f"[dim]{plain['weight_err']:.5f}[/dim]", "[dim]x1.00[/dim]", f"[dim]{plain_fwd:8.3f} ms[/dim]" if plain_fwd else "-") entry = dict(quant_s=None, size_bytes=None, weight_err=None, err_ratio=None, fwd_ms=None, rank=rank) results[f"{dtype_id}@{rank}"] = entry progress.reset(task) phase("quantizing (svd)") try: layer, entry["quant_s"] = make_quantized_linear(weight, use_svd=True, svd_rank=rank, **cfg) entry["size_bytes"] = layer_storage_bytes(layer) def fwd_fn(current=layer): return current(x) phase("timing dequant and forward") with time_limit(config_timeout, row_label): deq_out = layer.sdnq_dequantizer( layer.weight, layer.scale, zero_point=getattr(layer, "zero_point", None), svd_up=getattr(layer, "svd_up", None), svd_down=getattr(layer, "svd_down", None), ).to(torch.float32) fwd_fn() torch_device_module.synchronize() entry["weight_err"] = rel_err(deq_out, weight_fp32) del deq_out plain_err = plain.get("weight_err") entry["err_ratio"] = entry["weight_err"] / plain_err if plain_err else None entry["fwd_ms"] = bench(fwd_fn, warmup, iters) table.add_row(row_label, quant_cell(entry["quant_s"]), size_cell(entry["size_bytes"]), err_cell(entry["weight_err"]), f"x{entry['err_ratio']:.2f}" if entry["err_ratio"] else "-", f"{entry['fwd_ms']:8.3f} ms") del layer except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(row_label, quant_cell(entry["quant_s"]), "-", "-", "-", failure_text(e)) torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state live.update(panel) console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(panel) torch_device_module.empty_cache() report["dequant_svd_ranks"] = dict(shape=shape_label, results=results) return results def bench_hadamard_groups(shape_label, out_features, in_features, plain_results, selected_dtypes, iters, warmup, config_timeout=300): # the weight-side Hadamard group size setting, swept on int8 where rotation measured the # largest gain; unlike the attention slider it is not clamped to head dim if selected_dtypes is not None and "int8" not in selected_dtypes: return {} table = Table(box=box.SIMPLE_HEAVY) table.add_column("config") table.add_column("quant", justify="right") table.add_column("size", justify="right") table.add_column("weight err", justify="right") table.add_column("err vs plain", justify="right") table.add_column("fwd", justify="right") panel = Panel(table, title=f"hadamard group size sweep: int8 {shape_label} {dtype_label()}", subtitle="[dim]weight-side Hadamard group size setting; err vs plain below x1.00 = better reconstruction[/dim]", box=ROUNDED_BOX, expand=False) results = {} plain = plain_results.get("int8") or {} progress, task = live_progress() with Live(Group(panel, progress), console=console, refresh_per_second=4) as live: weight = make_source_weight(out_features, in_features) weight_fp32 = weight.to(torch.float32) generator = torch.Generator(device=torch_device).manual_seed(42) x = torch.randn(dequant_forward_tokens, in_features, device=torch_device, dtype=bench_dtype, generator=generator) if plain.get("weight_err") is not None: plain_fwd = plain.get("fwd_ms") table.add_row("[dim]int8[/dim]", f"[dim]{quant_cell(plain.get('quant_s'))}[/dim]", f"[dim]{size_cell(plain['size_bytes'])}[/dim]", f"[dim]{plain['weight_err']:.5f}[/dim]", "[dim]x1.00[/dim]", f"[dim]{plain_fwd:8.3f} ms[/dim]" if plain_fwd else "-") for index, hgroup in enumerate(hadamard_group_values, start=1): row_label = f"int8 + hadamard group {hgroup}" def phase(step, current_label=row_label, current_index=index): progress.update(task, description=f"hadamard group {current_index}/{len(hadamard_group_values)} {current_label}: {step}") entry = dict(quant_s=None, size_bytes=None, weight_err=None, err_ratio=None, fwd_ms=None, hgroup=hgroup) results[str(hgroup)] = entry progress.reset(task) phase("quantizing (hadamard)") try: layer, entry["quant_s"] = make_quantized_linear(weight, "int8", use_hadamard=True, hadamard_group_size=hgroup) entry["size_bytes"] = layer_storage_bytes(layer) def fwd_fn(current=layer): return current(x) phase("timing dequant and forward") with time_limit(config_timeout, row_label): deq_out = layer.sdnq_dequantizer( layer.weight, layer.scale, zero_point=getattr(layer, "zero_point", None), svd_up=getattr(layer, "svd_up", None), svd_down=getattr(layer, "svd_down", None), ).to(torch.float32) fwd_fn() torch_device_module.synchronize() entry["weight_err"] = rel_err(deq_out, weight_fp32) del deq_out plain_err = plain.get("weight_err") entry["err_ratio"] = entry["weight_err"] / plain_err if plain_err else None entry["fwd_ms"] = bench(fwd_fn, warmup, iters) table.add_row(row_label, quant_cell(entry["quant_s"]), size_cell(entry["size_bytes"]), err_cell(entry["weight_err"]), f"x{entry['err_ratio']:.2f}" if entry["err_ratio"] else "-", f"{entry['fwd_ms']:8.3f} ms") del layer except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(row_label, quant_cell(entry["quant_s"]), "-", "-", "-", failure_text(e)) torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state live.update(panel) console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(panel) torch_device_module.empty_cache() report["dequant_hadamard_groups"] = dict(shape=shape_label, results=results) return results def bench_quant_toggles(shape_label, out_features, in_features, plain_results, selected_dtypes, iters, warmup, config_timeout=300): # remaining checkbox-level settings: Dequantize using full precision off (scales kept in # the model dtype instead of fp32), Dynamic quantization (per-layer dtype escalation until # the loss threshold passes), and Quantize using GPU (quantization wall time on cpu) fp32_dtypes = [(dtype_id, label, cfg) for dtype_id, label, cfg in dequant_dtype_configs if dtype_id in toggle_dtypes and (selected_dtypes is None or dtype_id in selected_dtypes)] results = dict(fp32_off={}, dynamic={}, cpu_quant_s=None, gpu_quant_s=None) table = Table(box=box.SIMPLE_HEAVY) table.add_column("config") table.add_column("quant", justify="right") table.add_column("size", justify="right") table.add_column("weight err", justify="right") table.add_column("fwd", justify="right") panel = Panel(table, title=f"dequantize full precision off: {shape_label} {dtype_label()}", subtitle="[dim]Dequantize using full precision unchecked; full-precision rows repeated dimmed[/dim]", box=ROUNDED_BOX, expand=False) progress, task = live_progress() with Live(Group(panel, progress), console=console, refresh_per_second=4) as live: weight = make_source_weight(out_features, in_features) weight_fp32 = weight.to(torch.float32) generator = torch.Generator(device=torch_device).manual_seed(42) x = torch.randn(dequant_forward_tokens, in_features, device=torch_device, dtype=bench_dtype, generator=generator) for index, (dtype_id, label, cfg) in enumerate(fp32_dtypes, start=1): row_label = f"{label} + fp32 dequant off" def phase(step, current_label=row_label, current_index=index): progress.update(task, description=f"toggle {current_index}/{len(fp32_dtypes)} {current_label}: {step}") plain = plain_results.get(dtype_id) or {} if plain.get("weight_err") is not None: plain_fwd = plain.get("fwd_ms") table.add_row(f"[dim]{label}[/dim]", f"[dim]{quant_cell(plain.get('quant_s'))}[/dim]", f"[dim]{size_cell(plain['size_bytes'])}[/dim]", f"[dim]{plain['weight_err']:.5f}[/dim]", f"[dim]{plain_fwd:8.3f} ms[/dim]" if plain_fwd else "-") entry = dict(quant_s=None, size_bytes=None, weight_err=None, fwd_ms=None) results["fp32_off"][dtype_id] = entry progress.reset(task) phase("quantizing") try: layer, entry["quant_s"] = make_quantized_linear(weight, dequantize_fp32=False, **cfg) entry["size_bytes"] = layer_storage_bytes(layer) def fwd_fn(current=layer): return current(x) phase("timing dequant and forward") with time_limit(config_timeout, row_label): deq_out = layer.sdnq_dequantizer( layer.weight, layer.scale, zero_point=getattr(layer, "zero_point", None), svd_up=getattr(layer, "svd_up", None), svd_down=getattr(layer, "svd_down", None), ).to(torch.float32) fwd_fn() torch_device_module.synchronize() entry["weight_err"] = rel_err(deq_out, weight_fp32) del deq_out entry["fwd_ms"] = bench(fwd_fn, warmup, iters) table.add_row(row_label, quant_cell(entry["quant_s"]), size_cell(entry["size_bytes"]), err_cell(entry["weight_err"]), f"{entry['fwd_ms']:8.3f} ms") del layer except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(row_label, quant_cell(entry["quant_s"]), "-", "-", failure_text(e)) torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state live.update(panel) console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(panel) dyn_table = Table(box=box.SIMPLE_HEAVY) dyn_table.add_column("requested") dyn_table.add_column("chosen", justify="center") dyn_table.add_column("quant", justify="right") dyn_table.add_column("size", justify="right") dyn_table.add_column("weight err", justify="right") dyn_panel = Panel(dyn_table, title=f"dynamic quantization: {shape_label} {dtype_label()}", subtitle="[dim]Use Dynamic quantization: escalates to wider dtypes until normalized mse passes the loss threshold (default 10^-(bits/2))[/dim]", box=ROUNDED_BOX, expand=False) progress, task = live_progress() with Live(Group(dyn_panel, progress), console=console, refresh_per_second=4) as live: weight = make_source_weight(out_features, in_features) weight_fp32 = weight.to(torch.float32) for index, requested in enumerate(dynamic_quant_requests, start=1): progress.reset(task) progress.update(task, description=f"dynamic {index}/{len(dynamic_quant_requests)} requested {requested}") entry = dict(quant_s=None, chosen=None, size_bytes=None, weight_err=None) results["dynamic"][requested] = entry try: layer, entry["quant_s"] = make_quantized_linear(weight, requested, use_dynamic_quantization=True) entry["chosen"] = layer.sdnq_dequantizer.weights_dtype entry["size_bytes"] = layer_storage_bytes(layer) deq_out = layer.sdnq_dequantizer( layer.weight, layer.scale, zero_point=getattr(layer, "zero_point", None), svd_up=getattr(layer, "svd_up", None), svd_down=getattr(layer, "svd_down", None), skip_compile=True, ).to(torch.float32) entry["weight_err"] = rel_err(deq_out, weight_fp32) del deq_out, layer chosen_cell = entry["chosen"] if entry["chosen"] == requested else f"[yellow]{entry['chosen']}[/yellow]" dyn_table.add_row(requested, chosen_cell, quant_cell(entry["quant_s"]), size_cell(entry["size_bytes"]), err_cell(entry["weight_err"])) except Exception as e: entry["error"] = error_summary(e, 200) dyn_table.add_row(requested, "[red]unquantized[/red]", "-", "-", failure_text(e)) live.update(dyn_panel) console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(dyn_panel) # Quantize using GPU: one int8 layer quantized on cpu for the wall-time comparison try: weight_cpu = make_source_weight(out_features, in_features).cpu() _cpu_layer, results["cpu_quant_s"] = make_quantized_linear(weight_cpu, "int8", device="cpu") del _cpu_layer, weight_cpu results["gpu_quant_s"] = (plain_results.get("int8") or {}).get("quant_s") except Exception as e: results["cpu_quant_error"] = error_summary(e, 200) torch_device_module.empty_cache() report["dequant_toggles"] = dict(shape=shape_label, results=results) return results def make_source_conv_weight(out_channels, in_channels, kernel, seed=1234): # conv analogue of make_source_weight: a few high-magnitude input channels generator = torch.Generator(device=torch_device).manual_seed(seed) weight = torch.randn(out_channels, in_channels, kernel, kernel, device=torch_device, dtype=bench_dtype, generator=generator) * 0.02 weight[:, [1, in_channels // 2, in_channels - 2]] *= 8.0 return weight def fp32_conv_reference(x, weight_fp32, padding): # convs route through cudnn, which has its own tf32 switch on top of the matmul one tf32_matmul = torch.backends.cuda.matmul.allow_tf32 tf32_cudnn = torch.backends.cudnn.allow_tf32 torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False try: return torch.nn.functional.conv2d(x.to(torch.float32), weight_fp32, padding=padding) finally: torch.backends.cuda.matmul.allow_tf32 = tf32_matmul torch.backends.cudnn.allow_tf32 = tf32_cudnn def make_quantized_conv(weight, weights_dtype, use_quantized_matmul=False): from sdnq import SDNQConfig from sdnq.quantizer import sdnq_quantize_layer out_channels, in_channels, kh, kw = weight.shape conv = torch.nn.Conv2d(in_channels, out_channels, (kh, kw), padding=(kh // 2, kw // 2), bias=False, device=torch_device, dtype=bench_dtype) with torch.no_grad(): conv.weight.copy_(weight) config = SDNQConfig(weights_dtype=weights_dtype, quant_conv=True, use_quantized_matmul_conv=use_quantized_matmul, add_skip_keys=False) torch_device_module.synchronize() started = time.perf_counter() layer, _config = sdnq_quantize_layer(conv, config, torch_dtype=bench_dtype, param_name="bench.weight") torch_device_module.synchronize() quant_seconds = time.perf_counter() - started if not hasattr(layer, "sdnq_dequantizer"): raise RuntimeError(f"sdnq did not quantize the conv layer to {weights_dtype}") return layer, quant_seconds def bench_conv_section(iters, warmup, config_timeout=300): # the Quantize convolutional layers and Use quantized MatMul with conv settings, measured # on real Conv2d layers; dit models have no convs, this is for unet and vae model classes all_results = {} for shape_label, out_channels, in_channels, kernel, px in conv_shapes: table = Table(box=box.SIMPLE_HEAVY) table.add_column("config") table.add_column("quant", justify="right") table.add_column("size", justify="right") table.add_column("weight err", justify="right") table.add_column("fwd", justify="right") table.add_column("out err", justify="right") table.add_column("speedup", justify="right") panel = Panel(table, title=f"conv quantization: {shape_label} {dtype_label()}", subtitle="[dim]bias-free conv2d, batch 1, errors vs true fp32 (tf32 off); mm rows use the conv quantized matmul path[/dim]", box=ROUNDED_BOX, expand=False) results = {} base_ms = None progress, task = live_progress() with Live(Group(panel, progress), console=console, refresh_per_second=4) as live: progress.update(task, description=f"{shape_label}: preparing weights and fp32 reference") weight = make_source_conv_weight(out_channels, in_channels, kernel) weight_fp32 = weight.to(torch.float32) generator = torch.Generator(device=torch_device).manual_seed(42) x = torch.randn(1, in_channels, px, px, device=torch_device, dtype=bench_dtype, generator=generator) ref_out = fp32_conv_reference(x, weight_fp32, kernel // 2) for index, (config_id, weights_cfg, use_mm) in enumerate(conv_configs, start=1): label = f"{dtype_label()} conv2d" if weights_cfg is None else f"{weights_cfg['weights_dtype']}{' + conv mm' if use_mm else ''}" def phase(step, current_label=label, current_index=index, current_progress=progress, current_task=task, current_shape=shape_label): current_progress.update(current_task, description=f"conv {current_index}/{len(conv_configs)} {current_shape} {current_label}: {step}") entry = dict(quant_s=None, size_bytes=None, weight_err=None, fwd_ms=None, out_err=None) results[config_id] = entry progress.reset(task) try: if weights_cfg is None: layer = torch.nn.Conv2d(in_channels, out_channels, (kernel, kernel), padding=(kernel // 2, kernel // 2), bias=False, device=torch_device, dtype=bench_dtype) with torch.no_grad(): layer.weight.copy_(weight) entry["size_bytes"] = layer.weight.numel() * layer.weight.element_size() weight_err_cell = "-" else: phase("quantizing") layer, entry["quant_s"] = make_quantized_conv(weight, use_quantized_matmul=use_mm, **weights_cfg) entry["size_bytes"] = layer_storage_bytes(layer) deq_out = layer.sdnq_dequantizer( layer.weight, layer.scale, zero_point=getattr(layer, "zero_point", None), svd_up=getattr(layer, "svd_up", None), svd_down=getattr(layer, "svd_down", None), skip_quantized_matmul=use_mm, skip_compile=True, ).to(torch.float32) entry["weight_err"] = rel_err(deq_out, weight_fp32) weight_err_cell = err_cell(entry["weight_err"]) del deq_out def fwd_fn(current=layer, current_x=x): return current(current_x) phase("timing conv forward") with time_limit(config_timeout, label): fwd_fn() torch_device_module.synchronize() entry["fwd_ms"], entry["fwd_ms_sigma"] = bench_stats(fwd_fn, warmup, iters) entry["out_err"] = rel_err(fwd_fn(), ref_out) if base_ms is None: base_ms = entry["fwd_ms"] table.add_row(label, quant_cell(entry["quant_s"]), size_cell(entry["size_bytes"]), weight_err_cell, f"{entry['fwd_ms']:8.3f} ms", err_cell(entry["out_err"]), speedup_cell(base_ms, entry["fwd_ms"])) del layer except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(label, quant_cell(entry["quant_s"]), "-", "-", failure_text(e), "-", "-") torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state live.update(panel) console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(panel) torch_device_module.empty_cache() all_results[shape_label] = results report["dequant_conv"] = all_results return all_results def block_label(weights_cfg, use_mm, attention_spec): if weights_cfg is None: weights_part = dtype_label() else: weights_part = weights_cfg["weights_dtype"] + (" g16" if weights_cfg.get("group_size") == 16 else "") if use_mm: weights_part += " mm" return f"{weights_part} + {attention_spec}" def bench_block_section(iters, warmup, config_timeout=None, selected=None, geometries=None): global block_geometry # pylint: disable=global-statement all_results = {} for family, geometry in block_geometries.items(): if geometries is not None and family not in geometries: continue block_geometry = geometry geometry_iters = geometry.get("iters", iters) geometry_warmup = geometry.get("warmup", warmup) geometry_timeout = resolve_timeout(config_timeout, geometry.get("config_timeout")) results = bench_block_geometry(geometry_iters, geometry_warmup, config_timeout=geometry_timeout, selected=selected) all_results[family] = results report.setdefault("blocks", {})[family] = dict(geometry=dict(geometry), results=results, split=None) if not all_results: return {} # the first family run also lands at the flat block key, which replays and the buyback # veto fall back to when no family matches the reference shape primary = next(iter(all_results)) report["block"] = report["blocks"][primary] return all_results.get(primary, {}) def bench_block_geometry(iters, warmup, config_timeout=300, selected=None): # each row is a complete configuration measured end to end through a dit block, because # component speedups and errors do not compose multiplicatively configs = [c for c in block_configs if selected is None or c[0] in selected] sage_missing = {"sage": sage_attention() is None, "sage fp16 accum": sage_attention_fp16_accum() is None} accum_missing = not atten_supports_fp16_accum() def spec_unavailable(spec): if sage_missing.get(spec, False): return True spec_kwargs = block_attention_specs.get(spec) return accum_missing and isinstance(spec_kwargs, dict) and bool(spec_kwargs.get("use_fp16_accum")) skipped = [config_id for config_id, _w, _mm, spec in configs if spec_unavailable(spec)] if skipped: configs = [c for c in configs if c[0] not in skipped] emit(f"[dim]block: skipping {', '.join(skipped)}, sageattention or the sdnq accumulation kwarg is unavailable here[/dim]") if not configs: return {} hidden, heads, mlp_dim, tokens = block_geometry["hidden"], block_geometry["heads"], block_geometry["mlp_dim"], block_geometry["tokens"] table = Table(box=box.SIMPLE_HEAVY) table.add_column("config") table.add_column("size", justify="right") table.add_column("block ms", justify="right") table.add_column("speedup", justify="right") table.add_column("out err", justify="right") table.add_column("max tok err", justify="right") table.add_column("err x4 blocks", justify="right") panel = Panel(table, title=f"combined block: hidden={hidden} heads={heads} mlp={mlp_dim} tokens={tokens} {dtype_label()}", subtitle="[dim]dit block, fused qkv attention + mlp with residuals; err vs an fp32 reference block, max tok = worst single token, x4 = four stacked blocks[/dim]", box=ROUNDED_BOX, expand=False) def run_depth(block, x0, depth): h = x0 for _ in range(depth): h = block(h) return h results = {} base_ms = None progress, task = live_progress() with Live(Group(panel, progress), console=console, refresh_per_second=4) as live: progress.update(task, description="block: building master weights and the fp32 reference") master = make_block_master() generator = torch.Generator(device=torch_device).manual_seed(7) x = torch.randn(1, tokens, hidden, device=torch_device, dtype=bench_dtype, generator=generator) ref_block = build_block_module(dtype=torch.float32) ref_block.load_state_dict(master) ref_block.eval() def ref_attention(q, k, v): return torch.nn.functional.scaled_dot_product_attention(q, k, v) ref_block.attention_fn = ref_attention tf32_matmul = torch.backends.cuda.matmul.allow_tf32 tf32_cudnn = torch.backends.cudnn.allow_tf32 torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False try: with torch.no_grad(): x_fp32 = x.to(torch.float32) ref_out = ref_block(x_fp32) ref_out4 = run_depth(ref_block, x_fp32, 4) finally: torch.backends.cuda.matmul.allow_tf32 = tf32_matmul torch.backends.cudnn.allow_tf32 = tf32_cudnn del ref_block, x_fp32 torch_device_module.empty_cache() for index, (config_id, weights_cfg, use_mm, attention_spec) in enumerate(configs, start=1): label = block_label(weights_cfg, use_mm, attention_spec) def phase(step, current_label=label, current_index=index): progress.update(task, description=f"block config {current_index}/{len(configs)} {current_label}: {step}") entry = dict(ms=None, err=None, max_err=None, err4=None, size_bytes=None, label=label) results[config_id] = entry progress.reset(task) phase("quantizing block") try: block = build_bench_block(master, weights_cfg, use_mm, attention_spec) entry["size_bytes"] = block_storage_bytes(block) def fn(current=block): with torch.no_grad(): return current(x) phase("compiling") with time_limit(config_timeout, label): out = fn() torch_device_module.synchronize() entry["ms"], entry["ms_sigma"] = bench_stats(fn, warmup, iters, on_phase=phase) entry["err"] = rel_err(out, ref_out) entry["max_err"] = max_token_err(out, ref_out) phase("measuring depth-4 error") with torch.no_grad(): entry["err4"] = rel_err(run_depth(block, x, 4), ref_out4) phase("timing identity-attention variant") real_attention_fn = block.attention_fn def identity_attention_fn(q, k, v): # pylint: disable=unused-argument # same shapes and permutes, zero attention flops return v block.attention_fn = identity_attention_fn try: entry["identity_ms"], entry["identity_ms_sigma"] = bench_stats(fn, warmup, iters, on_phase=phase) finally: block.attention_fn = real_attention_fn del block, out if base_ms is None: base_ms = entry["ms"] table.add_row(label, size_cell(entry["size_bytes"]), f"{entry['ms']:8.3f} ms", speedup_cell(base_ms, entry["ms"]), err_cell(entry["err"]), err_cell(entry["max_err"]), err_cell(entry["err4"])) except Exception as e: entry["error"] = error_summary(e, 200) table.add_row(label, size_cell(entry["size_bytes"]) if entry["size_bytes"] else "-", "-", "-", "-", "-", failure_text(e)) torch._dynamo.reset() # pylint: disable=protected-access # drop failed compile state live.update(panel) console.line() # separate sections; Live's final frame also lacks a trailing newline when output is piped transcript.append(panel) torch_device_module.empty_cache() notes = [] sized = [(entry["label"], entry["ms"], entry["err"]) for entry in results.values() if entry.get("ms") and entry.get("err") is not None and math.isfinite(entry["err"])] frontier = sorted((c for c in sized if not any(o is not c and o[1] <= c[1] and o[2] <= c[2] for o in sized)), key=lambda c: c[1]) if len(frontier) > 1: notes.append("speed/error frontier: " + ", ".join(f"{label} {ms:.2f}ms err {err:.4f}" for label, ms, err in frontier)) growth = [entry["err4"] / entry["err"] for entry in results.values() if entry.get("err") and entry.get("err4") and math.isfinite(entry["err4"] / entry["err"])] if growth: notes.append(f"four stacked blocks: error grows x{sum(growth) / len(growth):.2f} avg, residuals dampen compounding") weights_mode = str(getattr(shared.opts, "sdnq_quantize_weights_mode", "")) current_id = None if weights_mode == "int8" and getattr(shared.opts, "sdnq_quantize_matmul_mode", "disabled") != "disabled": current_id = "int8-mm-atten" if "SDNQ attention" in shared.opts.cross_attention_optimization else "int8-mm" if current_id and results.get(current_id, {}).get("ms"): notes.append(f"current config runs the {results[current_id]['label']} row for int8-quantized models") if any(entry.get("ms") for config_id, entry in results.items() if config_id.endswith("sagefp16")): notes.append("sage fp16 accum: benchmark-only for sage (sdnext pins its fp32 accum); the equivalent sdnq mode is the FP16 Accumulation checkbox with pv matmul disabled, same overflow tail") if notes: emit(Panel("\n".join(notes), title="block notes", box=ROUNDED_BOX)) return results def emit_block_splits(): # instrument B reads the attention tables, so the split renders once both sections are in for family, data in (report.get("blocks") or {}).items(): data["split"] = block_split_table(family, data["geometry"], data["results"]) def block_split_table(family, geometry, results): # compute split per config from two independent instruments: A subtracts the identity- # attention variant timed inside the block, B reads the standalone attention table at the # same geometry from this run; a speedup ceiling is only stated where the two agree head_dim = geometry.get("head_dim") or geometry["hidden"] // geometry["heads"] expected_geometry = f"batch=1 heads={geometry['heads']} tokens={geometry['tokens']} head_dim={head_dim}" attention_preset = None for preset_name, data in (report.get("attention") or {}).items(): if data.get("geometry") == expected_geometry and not shape_presets.get(preset_name, {}).get("mask_fn"): attention_preset = preset_name break attention_results = (report.get("attention") or {}).get(attention_preset, {}).get("results", {}) if attention_preset else {} spec_by_config = {config_id: spec for config_id, _w, _mm, spec in block_configs} budgets = (0.5, 0.3, 0.15) table = Table(box=box.SIMPLE_HEAVY) table.add_column("config") table.add_column("block ms", justify="right") table.add_column("rest ms", justify="right") table.add_column("attn A", justify="right") table.add_column("attn B", justify="right") table.add_column("agree", justify="right") for budget in budgets: table.add_column(f"ceil@{int(budget * 100)}%", justify="right") split = {} for config_id, entry in results.items(): ms, identity_ms = entry.get("ms"), entry.get("identity_ms") if not ms or not identity_ms: continue attn_a = ms - identity_ms if attn_a <= 0: continue attention_id = block_spec_attention_ids.get(spec_by_config.get(config_id)) attention_entry = attention_results.get(attention_id) or {} attn_b = attention_entry.get("ms") agree = None if attn_b: # subtraction amplifies the relative sigma of instrument A by ms/attn_a sigma_a = (row_sigma(entry) or 0.0) * (ms / attn_a) sigma_b = row_sigma(attention_entry, "ms") or 0.0 threshold = max(verdict_z * math.sqrt(sigma_a * sigma_a + sigma_b * sigma_b), run_drift_sigma(), split_instrument_offset) agree = abs(math.log(attn_a / attn_b)) <= threshold ceilings = {budget: ms / (budget * attn_a + identity_ms) for budget in budgets} if agree else None split[config_id] = dict(ms=ms, rest_ms=identity_ms, attn_a_ms=attn_a, attn_b_ms=attn_b, agree=agree, ceilings=ceilings) agree_cell = "-" if agree is None else ("yes" if agree else "[yellow]no[/yellow]") ceiling_cells = [f"{ceilings[budget]:.2f}x" if ceilings else "-" for budget in budgets] table.add_row(entry["label"], f"{ms:8.3f} ms", f"{identity_ms:8.3f} ms", f"{attn_a:8.3f} ms", f"{attn_b:8.3f} ms" if attn_b else "-", agree_cell, *ceiling_cells) if split: subtitle = "[dim]rest = identity-attention variant; A = block minus rest, B = standalone attention table" subtitle += f" ({attention_preset})" if attention_preset else " (no matching attention preset this run)" subtitle += "; ceilings are per-block upper bounds at the given kv budget, generation adds te/vae/projections[/dim]" emit(Panel(table, title=f"compute split: {family}", subtitle=subtitle, box=ROUNDED_BOX, expand=False)) disagreements = [config_id for config_id, row in split.items() if row["agree"] is False] if disagreements: emit(f"[yellow]split instruments disagree on {', '.join(disagreements)}; ceilings withheld there, treat the split with suspicion[/yellow]") return split def measured(results, config_id): entry = results.get(config_id) or {} ms = entry.get("ms") return (ms, entry.get("err")) if ms is not None else (None, None) def best_config(results): # lowest error among rows within 5% of the fastest sdnq time candidates = [(config_id, entry["ms"], entry["err"]) for config_id, entry in results.items() if entry.get("ms") is not None and config_id not in external_config_ids and config_id not in sparse_config_ids and config_id not in unsafe_config_ids] if not candidates: return None fastest = min(ms for _config_id, ms, _err in candidates) near_fastest = [(config_id, ms, err) for config_id, ms, err in candidates if ms <= fastest * 1.05] return min(near_fastest, key=lambda item: item[2])[0] def config_settings(kwargs): # ui settings tuple a bench config corresponds to; None for the external baselines # and the unquantized row, which are not reachable states of the attention dropdowns if not kwargs or not kwargs.get("do_quantize", True): return None matmul = kwargs.get("matmul_dtype", "auto") pv = kwargs.get("pv_matmul_dtype", "auto") return { "matmul": "enabled" if matmul == "auto" else matmul, "pv": "disabled" if pv == "auto" else pv, "smooth": bool(kwargs.get("smooth_k", False)), "hadamard": bool(kwargs.get("use_hadamard", False)), "accum": bool(kwargs.get("use_fp16_accum", False)), } def select_attention_config(results): # the attention settings jointly pick one kernel config, so they are judged as one: # rank the measured quantized rows by the same rule as the per-shape star instead of # testing each setting alone against int8, which can assemble a settings tuple no # row ever measured. rows breaching the incremental error cap vs int8 are held out # of the ranking but kept for citation int8_err_ref = measured(results, "int8")[1] pool, capped = [], [] for config_id, label, kwargs in bench_configs: settings = config_settings(kwargs) if settings is None or config_id in sparse_config_ids: # a sparse row shares a settings tuple with its dense row but is a stage over it, not a setting continue if settings["accum"] and settings["pv"] == "disabled": continue # the unsafe accumulation combo is never a candidate; the accum row cites it directly ms, err = measured(results, config_id) if ms is None: continue entry = dict(config_id=config_id, label=label.removeprefix("sdnq "), ms=ms, err=err, settings=settings) if err and int8_err_ref and err > int8_err_ref * recommend_error_cap: capped.append(entry) else: pool.append(entry) if not pool: return None, pool, capped fastest = min(entry["ms"] for entry in pool) window = [entry for entry in pool if entry["ms"] <= fastest * 1.05] chosen = min(window, key=lambda entry: entry["err"] if entry["err"] is not None else float("inf")) return chosen, pool, capped def best_in_subset(entries): # the star rule inside one setting's subset, for citing its strongest variant if not entries: return None fastest = min(entry["ms"] for entry in entries) window = [entry for entry in entries if entry["ms"] <= fastest * 1.05] return min(window, key=lambda entry: entry["err"] if entry["err"] is not None else float("inf")) def build_recommendations(all_results, fp8_result, prep_status, block_results=None, block_variants=None): # prefer an image dit shape with the full config set as reference, then the video shapes reference = None for preset in recommendation_presets: if preset in all_results and measured(all_results[preset], "int8")[0] is not None: reference = preset break if reference is None: emit(f"[yellow]attention recommendations need a successful int8 run at a self-attention reference shape ({', '.join(recommendation_presets)}); none was benchmarked[/yellow]") return results = all_results[reference] base_ms, _base_err = measured(results, "base") noquant_ms, noquant_err = measured(results, "noquant") int8_ms, int8_err = measured(results, "int8") table = Table(title=f"recommended settings (Compute Settings -> SDNQ Attention), measured at the {reference} shape", box=ROUNDED_BOX) table.add_column("setting") table.add_column("current", justify="center") table.add_column("recommended", justify="center") table.add_column("reason") def current(key): return str(getattr(shared.opts, key)) rows = [] # the settings below jointly pick one kernel config, so the verdict is joint: select # the best measured quantized row (star rule with the error cap), then gate that one # config against the unquantized sdnq row: quantization has to pay for its own prep # and clear the shared speed margin, at this run's own measured noise level and at a # z adjusted for having selected the best of the pool; too close to call keeps the # current settings chosen, pool, capped = select_attention_config(results) gate = None quant_reason = "no quantized attention config ran at this shape" if chosen is not None and noquant_ms: gate_sigma = pair_sigma(results.get(chosen["config_id"]) or {}, "ms", results.get("noquant") or {}, "ms") gate = speed_verdict(chosen["ms"], noquant_ms, sigma=gate_sigma, z=sidak_z_for(len(pool))) if gate == "faster": quant_reason = f"{chosen['label']} measured x{noquant_ms / chosen['ms']:.2f} vs unquantized sdnq attention" if chosen["err"] and noquant_err: quant_reason += f", error {chosen['err']:.5f} vs {noquant_err:.5f}" elif gate == "inconclusive": quant_reason = f"best quantized config ({chosen['label']}) measured x{noquant_ms / chosen['ms']:.2f} vs unquantized sdnq attention, too close to the x{1 / recommend_speed_margin:.2f} margin to call at this run's noise (ratio sigma {gate_sigma:.1%}); keeping the current setting" elif chosen["ms"] < noquant_ms: quant_reason = f"best quantized config ({chosen['label']}) measured only x{noquant_ms / chosen['ms']:.2f} vs unquantized sdnq attention, under the x{1 / recommend_speed_margin:.2f} margin the verdict requires" if chosen["err"] and noquant_err: quant_reason += f"; unquantized keeps error {noquant_err:.5f} vs {chosen['err']:.5f}" else: quant_reason = f"unquantized sdnq measured x{chosen['ms'] / noquant_ms:.2f} vs the best quantized config ({chosen['label']}) with lower error; quantization prep outweighs the kernel gain on this gpu" elif chosen is not None and base_ms: gate = "faster" if base_ms / chosen["ms"] >= 1.10 else "not_faster" quant_reason = f"{chosen['label']} measured x{base_ms / chosen['ms']:.2f} vs torch sdpa; unquantized sdnq row unavailable" use_quantized = gate == "faster" qk_inconclusive = gate == "inconclusive" # buyback costs are judged at the block geometry matching the reference family when # one was measured, so a krea2 verdict uses the krea2-width block block_rows = block_results or {} block_geo = (report.get("block") or {}).get("geometry") or {} if block_variants: block_family = reference if reference in block_variants else next(iter(block_variants)) variant = block_variants.get(block_family) or {} if variant.get("results"): block_rows = variant["results"] block_geo = variant.get("geometry") or {} def block_ms(config_id): return (block_rows.get(config_id) or {}).get("ms") def buyback_cost(kernel_config_id, block_config_id): # cost of a prep toggle, judged at the most representative measured scope: the # kernel rows hand the prep contiguous q/k/v, real models hand it strided views # from the fused qkv projection, which the block section measures toggle_ms, toggle_err = measured(results, kernel_config_id) if not (toggle_ms and int8_ms and int8_err and toggle_err): return None, None, None gain = int8_err / toggle_err if toggle_err > 0 else 1.0 kernel_cost = toggle_ms / int8_ms - 1.0 block_base_ms, block_toggle_ms = block_ms("int8-mm-atten"), block_ms(block_config_id) if block_base_ms and block_toggle_ms: cost = block_toggle_ms / block_base_ms - 1.0 geo = f" ({block_geo['hidden']}-wide, hd{block_geo['hidden'] // block_geo['heads']})" if block_geo.get("hidden") and block_geo.get("heads") else "" return gain, cost, f"{cost:+.0%} at block scope{geo} with strided qkv ({kernel_cost:+.0%} on contiguous kernel tensors)" return gain, kernel_cost, f"{kernel_cost:+.0%} time" smooth_gain, smooth_cost, smooth_cost_note = buyback_cost("smooth", "int8-mm-smooth") smooth_rec = smooth_reason = None if smooth_gain is not None: smooth_rec = smooth_gain >= 1.3 and smooth_cost <= 0.25 # attention-level cost is a few percent end to end smooth_reason = f"int8 error x{smooth_gain:.1f} lower for {smooth_cost_note}" hadamard_gain, hadamard_cost, hadamard_cost_note = buyback_cost("hadamard", "int8-mm-hadamard") hadamard_rec = hadamard_reason = None if hadamard_gain is not None: hadamard_rec = hadamard_gain >= 1.3 and hadamard_cost <= 0.15 hadamard_reason = f"int8 error x{hadamard_gain:.1f} lower for {hadamard_cost_note}; hangs torch compile on non pow2 head dims (SD 1.5)" # rows decompose the one selected config; each reason cites the measured row that # isolates its own setting where one was benchmarked by_settings = {} for entry in pool + capped: s = entry["settings"] by_settings[(s["matmul"], s["pv"], s["smooth"], s["hadamard"], s["accum"])] = entry def sibling(entry, **overrides): s = dict(entry["settings"], **overrides) other = by_settings.get((s["matmul"], s["pv"], s["smooth"], s["hadamard"], s["accum"])) return None if other is entry else other def block_buyback_survives(block_config_id): # None when the block section did not measure the pair; False when the toggle # left block output error unchanged, so its kernel-scope buyback is cosmetic here variant_err = (block_rows.get(block_config_id) or {}).get("err") base_block_err = (block_rows.get("int8-mm-atten") or {}).get("err") if not (variant_err and base_block_err): return None return variant_err < base_block_err * 0.98 fp8qk_ms, fp8qk_err = measured(results, "fp8qk") if use_quantized: qk_choice = chosen["settings"]["matmul"] qk_reason = quant_reason faster_capped = [entry for entry in capped if entry["ms"] < chosen["ms"]] if faster_capped and int8_err: fastest_capped = min(faster_capped, key=lambda entry: entry["ms"]) qk_reason += f"; {fastest_capped['label']} is x{chosen['ms'] / fastest_capped['ms']:.2f} faster but multiplies error x{fastest_capped['err'] / int8_err:.1f} over int8" if qk_choice == "enabled": qk_reason += "; enabled resolves to int8, uint8 remaps to int8" if fp8qk_ms and int8_ms and qk_choice != "float8_e4m3fn": qk_reason += f"; float8 qk measured x{int8_ms / fp8qk_ms:.2f} vs int8" if fp8qk_err and int8_err: qk_reason += f", error {fp8qk_err:.5f} vs {int8_err:.5f}" elif fp8_result["qk"][0] and not fp8qk_ms: qk_reason += "; float8 compiles here but was not benchmarked at this shape" elif qk_inconclusive: qk_choice = current("sdnq_attention_matmul_type") qk_reason = quant_reason else: qk_choice = "disabled" qk_reason = quant_reason rows.append(("MatMul type", current("sdnq_attention_matmul_type"), qk_choice, qk_reason)) pv_names = {"int8": "int8", "float16": "fp16", "float8_e4m3fn": "fp8"} if use_quantized: pv_choice = chosen["settings"]["pv"] if pv_choice != "disabled": without = sibling(chosen, pv="disabled") pv_reason = f"{pv_names[pv_choice]} pv rides the selected config" if without: pv_reason = f"{pv_names[pv_choice]} pv measured x{without['ms'] / chosen['ms']:.2f} over the same stack without pv" if chosen["err"] and without["err"]: pv_reason += f", error {chosen['err']:.5f} vs {without['err']:.5f}" else: best_pv = best_in_subset([entry for entry in pool + capped if entry["settings"]["pv"] != "disabled"]) pv_reason = "disabled keeps pv unquantized" if best_pv is not None: pv_label = f"{pv_names[best_pv['settings']['pv']]} pv ({best_pv['label']})" if best_pv["err"] and int8_err and best_pv["err"] > int8_err * recommend_error_cap: pv_reason = f"{pv_label} is x{chosen['ms'] / best_pv['ms']:.2f} the speed of the selected config but multiplies error x{best_pv['err'] / int8_err:.1f} over int8; disabled keeps pv unquantized" else: pv_reason += f"; {best_pv['label']} measured x{chosen['ms'] / best_pv['ms']:.2f} the speed of the selected config" if best_pv["err"] and chosen["err"]: pv_reason += f", error {best_pv['err']:.5f} vs {chosen['err']:.5f}" rows.append(("PV MatMul type", current("sdnq_attention_pv_matmul_type"), pv_choice, pv_reason)) elif qk_inconclusive: rows.append(("PV MatMul type", current("sdnq_attention_pv_matmul_type"), current("sdnq_attention_pv_matmul_type"), "the qk verdict above is inconclusive; pv follows it")) else: rows.append(("PV MatMul type", current("sdnq_attention_pv_matmul_type"), "disabled", "qk quantization is not recommended above; pv on an unquantized qk path was not measured")) # the accumulation flag only changes fp16 dots: scaled and overflow-proof on quantized # float16 types, inert on int8 types, and the sage-style unsafe mode on an unquantized # pv. only the scaled combos are candidates; the unsafe one is cited, never recommended accum_key = "sdnq_attention_use_fp16_accum" if atten_supports_fp16_accum(): accum_current = current(accum_key) if hasattr(shared.opts, accum_key) else "-" pvaccum_ms, _pvaccum_err = measured(results, "pvaccum") unsafe_note = "with an unquantized pv the flag is the sage-style unsafe accumulation instead (the stress rows demonstrate its overflow)" if pvaccum_ms and int8_ms: unsafe_note += f", measured x{int8_ms / pvaccum_ms:.2f} vs int8 qk" if use_quantized and chosen["settings"]["accum"]: without = sibling(chosen, accum=False) accum_reason = f"part of the selected config ({chosen['label']})" if without: accum_reason = f"scaled fp16 accumulation measured x{without['ms'] / chosen['ms']:.2f} over the same stack without it" if chosen["err"] and without["err"]: accum_reason += f", error {chosen['err']:.5f} vs {without['err']:.5f}" accum_reason += "; safe here because every fp16 dot in the stack is quantized and pre-scaled" rows.append(("Use FP16 Accumulation", accum_current, "True", accum_reason)) elif qk_inconclusive: rows.append(("Use FP16 Accumulation", accum_current, accum_current, "the qk verdict above is inconclusive; the accumulation flag follows it")) else: best_accum = best_in_subset([entry for entry in pool + capped if entry["settings"]["accum"]]) if chosen is not None and best_accum is not None: accum_reason = f"the strongest scaled-accum config ({best_accum['label']}) measured x{chosen['ms'] / best_accum['ms']:.2f} the speed of the selected config" else: accum_reason = "no scaled-accum config was measured at this shape" rows.append(("Use FP16 Accumulation", accum_current, "False", f"{accum_reason}; {unsafe_note}")) def toggle_row(setting_key, ui_name, on, buyback_reason, subset, block_config_id, caveat=None): # a toggle the selected config carries keeps the buyback evidence as its reason; # one it omits cites the strongest variant that carried it, plus the block-scope # cross-check when that shows the buyback never reached block output if on: reason = buyback_reason or f"part of the selected config ({chosen['label']})" if caveat and caveat not in reason: reason += f"; {caveat}" else: reason = "the selected config omits it" best_variant = best_in_subset(subset) if best_variant is not None: reason += f"; {best_variant['label']} measured x{chosen['ms'] / best_variant['ms']:.2f} the speed of the selected config" if best_variant["err"] and chosen["err"]: reason += f", error {best_variant['err']:.5f} vs {chosen['err']:.5f}" if block_buyback_survives(block_config_id) is False: reason += "; the block cross-check found no output-error change from it at this shape" rows.append((ui_name, current(setting_key), str(bool(on)), reason)) if use_quantized: toggle_row("sdnq_attention_smooth_k", "Use Smooth K", chosen["settings"]["smooth"], smooth_reason, [entry for entry in pool + capped if entry["settings"]["smooth"]], "int8-mm-smooth") toggle_row("sdnq_attention_use_hadamard", "Use Hadamard", chosen["settings"]["hadamard"], hadamard_reason, [entry for entry in pool + capped if entry["settings"]["hadamard"]], "int8-mm-hadamard", caveat="hangs torch compile on non pow2 head dims (SD 1.5)") else: if smooth_rec is not None: rows.append(("Use Smooth K", current("sdnq_attention_smooth_k"), str(smooth_rec), smooth_reason)) if hadamard_rec is not None: rows.append(("Use Hadamard", current("sdnq_attention_use_hadamard"), str(hadamard_rec), hadamard_reason)) rows.append(("Hadamard Group Size", current("sdnq_attention_hadamard_group_size"), "256", "values above head dim are clamped; non pow2 values floor to the nearest power of 2")) differing = 0 for setting, current_value, recommended, reason in rows: if current_value == recommended: table.add_row(setting, f"[green]{current_value}[/green]", recommended, reason) else: differing += 1 table.add_row(setting, f"[yellow]{current_value}[/yellow]", recommended, reason) emit(table) if differing == 1: emit("[yellow]1 setting differs from the recommended value, change it in Compute Settings -> SDNQ Attention[/yellow]") elif differing: emit(f"[yellow]{differing} settings differ from the recommended values, change them in Compute Settings -> SDNQ Attention[/yellow]") else: emit("[green]current settings already match the recommendations[/green]") if not use_quantized and not qk_inconclusive: emit("[dim]qk quantization is not worth it here, so MatMul type recommends disabled; the smooth k and hadamard rows show what to pick if it is enabled anyway[/dim]") notes = [] # per-shape joint verdicts: the table above judges one reference shape, but the # settings are global and workloads differ; a split gpu (video wins, image loses) # shows here rather than being averaged away, and each shape names the config its # star selected. cross and te shapes are prep-dominated, skip them shape_verdicts = [] for shape_label in recommendation_presets: shape_results = all_results.get(shape_label) if not shape_results: continue shape_noquant_ms, _err = measured(shape_results, "noquant") shape_chosen, shape_pool, _shape_capped = select_attention_config(shape_results) if shape_noquant_ms and shape_chosen is not None: shape_sigma = pair_sigma(shape_results.get(shape_chosen["config_id"]) or {}, "ms", shape_results.get("noquant") or {}, "ms") shape_test = speed_verdict(shape_chosen["ms"], shape_noquant_ms, sigma=shape_sigma, z=sidak_z_for(len(shape_pool))) word = {"faster": "enabled", "not_faster": "disabled", "inconclusive": "inconclusive"}[shape_test] shape_verdicts.append((word, f"{shape_label} {word} ({shape_chosen['label']} x{shape_noquant_ms / shape_chosen['ms']:.2f})")) if len(shape_verdicts) > 1: split = len({word for word, _text in shape_verdicts}) > 1 line = f"verdict by shape: {', '.join(text for _word, text in shape_verdicts)}" if split: notes.append(f"[yellow]{line}; settings are global, pick for the shapes you generate at[/yellow]") else: notes.append(line) labels = {config_id: config_label(config_id, label) for config_id, label, _kwargs in bench_configs} best_id = best_config(results) if base_ms and best_id is not None: best_ms = results[best_id]["ms"] if best_ms < base_ms * 0.95: notes.append(f"worth enabling: {labels[best_id]} measured x{base_ms / best_ms:.2f} vs torch sdpa at {reference}") else: notes.append(f"[yellow]not worth enabling: best is {labels[best_id]} at x{base_ms / best_ms:.2f} vs torch sdpa, the sdp override costs performance here[/yellow]") # the individual setting rows above compose into one config; cite it as actually measured full_ms, full_err = measured(results, "full") if full_ms and base_ms and int8_ms: stack_note = f"all options on (smooth k + hadamard + int8 pv over int8 qk): x{base_ms / full_ms:.2f} vs torch sdpa" if full_err is not None: stack_note += f", error {full_err:.5f}" if int8_err is not None: stack_note += f"; int8 qk alone x{base_ms / int8_ms:.2f}, error {int8_err:.5f}" notes.append(stack_note) if prep_status == "failing_dynamic": notes.append("[red]generation fails here: compile cannot build the dynamic-shape prep. fix: set SDNQ_COMPILE_KWARGS='{\"dynamic\": false}' (verified on this machine) or install msvc build tools; numbers above use that workaround[/red]") elif prep_status == "failing": notes.append("[red]generation fails here: torch compile is broken. fix: set SDNQ_USE_TORCH_COMPILE=0 or install msvc build tools; numbers above use eager prep[/red]") if not fp8_result["qk"][0]: if fp8_failure_is_capability(fp8_result["qk"][1]): notes.append("[red]float8_e4m3fn unsupported on this gpu: selecting it in either dropdown fails generation[/red]") else: notes.append("[red]float8_e4m3fn failed to compile here: selecting it in either dropdown fails generation[/red]; not the hardware signature, so a torch or triton issue rather than the gpu") sage_ms, _sage_err = measured(results, "sage") if sage_ms and int8_ms: notes.append(f"vs sage at {reference}: sdnq int8 {int8_ms:.2f} ms, sage {sage_ms:.2f} ms; sdnq also covers masks, gqa, causal") sagefp16_ms, _sagefp16_err = measured(results, "sagefp16") pvaccum_ms, _pvaccum_err = measured(results, "pvaccum") if sagefp16_ms and pvaccum_ms: notes.append(f"sage fp16 accum {sagefp16_ms:.2f} ms vs sdnq unquantized-pv accum {pvaccum_ms:.2f} ms: the same accumulation mode with the same overflow tail; sdnext reaches the sdnq one via SDNQ Attention use FP16 Accumulation with PV MatMul disabled, sage itself stays pinned to fp32 accum") elif sagefp16_ms: notes.append("sage fp16 accum: benchmark-only for sage (sdnext pins its fp32 accum on sm86); the equivalent sdnq mode is the unquantized-pv accum config") # compare flash against the config the verdict above actually recommends amdflash_ms, _amdflash_err = measured(results, "amdflash") if amdflash_ms: if use_quantized and int8_ms: notes.append(f"vs triton flash at {reference}: sdnq int8 {int8_ms:.2f} ms, flash {amdflash_ms:.2f} ms; sdnq also covers masks, its gain tracks this gpu's int8 throughput") elif noquant_ms: notes.append(f"vs triton flash at {reference}: unquantized sdnq {noquant_ms:.2f} ms, flash {amdflash_ms:.2f} ms; sdnq also covers masks") # cross-reference the combined block section when it ran in this invocation: kernel-scope # error buybacks can vanish at block output, and per-call costs differ inside a real block block_results = report.get("block", {}).get("results", {}) block_base = block_results.get("int8-mm-atten") or {} if block_base.get("err") and block_base.get("ms"): cross_parts = [] buyback_seen = False for variant_key, setting_name in (("int8-mm-smooth", "smooth k"), ("int8-mm-hadamard", "hadamard"), ("int8-mm-atten-full", "full stack (smooth + hadamard + int8 pv)")): variant_entry = block_results.get(variant_key) or {} if variant_entry.get("err") and variant_entry.get("ms"): cross_parts.append(f"{setting_name} output error {variant_entry['err']:.5f} vs {block_base['err']:.5f} at {variant_entry['ms'] - block_base['ms']:+.2f} ms per block") if variant_entry["err"] < block_base["err"] * 0.98: buyback_seen = True if cross_parts: cross_note = f"block-level cross-check vs plain int8 attention: {'; '.join(cross_parts)}" if not buyback_seen: cross_note += "; the kernel-scope error buybacks above did not change block output error at this image shape" notes.append(cross_note) notes.append("*: best config per shape, lowest error among rows within 5% of the fastest sdnq time") notes.append("default settings run the int8 qk row: MatMul type enabled resolves to int8, pv stays disabled") notes.append("prep: q/k/v quantization before the kernel, included in the median; shrinks where compile fuses it") notes.append("[yellow]non pow2 head dims (sd 1.5): quantized matmul fails compile, hadamard hangs it; disable quantized matmul there or set SDNQ_COMPILE_KWARGS='{\"dynamic\": false}'[/yellow]") notes.append("settings above apply on the next generation; the SDNQ attention SDP override needs a restart") notes.append("errors use synthetic outlier-heavy keys; real models, especially qk-normed dits, sit lower") emit(Panel("\n".join(notes), title="notes", box=ROUNDED_BOX)) report["recommendations_attention"] = dict(reference=reference, rows=[list(row) for row in rows], notes=notes) def ratio_text(numerator_ms, denominator_ms): if not numerator_ms or not denominator_ms: return None return f"x{numerator_ms / denominator_ms:.2f}" def build_dequant_recommendations(dequant_results, weight_dequant_result, variant_results=None, float_mm_results=None, sweep_results=None): reference = None for shape_label, _out_features, _in_features in dequant_shapes: results = dequant_results.get(shape_label) if results and results.get("int8", {}).get("eager_ms") is not None: reference = shape_label break if reference is None: emit("[yellow]no successful dequant benchmark, sdnq recommendations unavailable[/yellow]") return results = dequant_results[reference] base_ms = results.get("bf16", {}).get("fwd_ms") def entry(dtype_id): return results.get(dtype_id) or {} def current(key): return str(getattr(shared.opts, key)) table = Table(title=f"recommended settings (Compute Settings -> SDNQ), measured at the {reference} layer", box=ROUNDED_BOX) table.add_column("setting") table.add_column("current", justify="center") table.add_column("recommended", justify="center") table.add_column("reason") rows = [] weights_mode = str(getattr(shared.opts, "sdnq_quantize_weights_mode", "int8")) mode_ids = {cfg["weights_dtype"]: dtype_id for dtype_id, _label, cfg in dequant_dtype_configs} mm_id = mode_ids.get(weights_mode, "int8") # the compile toggle became automatic: sdnq compiles whenever triton is available and # SDNQ_USE_TORCH_COMPILE is the only override, so the verdict is a note addressed to the # env var rather than a settings row. it stays a symmetric faster/slower test (margin 1.0) # at the layer-forward scope compile gates, voted across every measured shape; the # standalone kernel can lose to eager on small layers from launch overhead alone, which # is why it never decides the verdict compile_id = mm_id if entry(mm_id).get("fwd_eager_ms") and entry(mm_id).get("fwd_compiled_ms") else "int8" fwd_votes = [] for shape_label in dequant_results: vote_row = (dequant_results.get(shape_label) or {}).get(compile_id) or {} fwd_eager, fwd_compiled = vote_row.get("fwd_eager_ms"), vote_row.get("fwd_compiled_ms") if fwd_eager and fwd_compiled: test = speed_verdict(fwd_compiled, fwd_eager, sigma=pair_sigma(vote_row, "fwd_compiled_ms", vote_row, "fwd_eager_ms"), margin=1.0) fwd_votes.append((fwd_eager / fwd_compiled, test)) int8_eager, int8_compiled = entry("int8").get("eager_ms"), entry("int8").get("compiled_ms") try: from sdnq.common import use_torch_compile as compile_active except Exception: compile_active = None compile_origin = "SDNQ_USE_TORCH_COMPILE" if os.environ.get("SDNQ_USE_TORCH_COMPILE", None) is not None else "automatic with triton" compile_state = {True: "on", False: "off"}.get(compile_active, "unknown") compile_note = None if fwd_votes: ratios = sorted(ratio for ratio, _test in fwd_votes) span = f"x{ratios[0]:.2f}" if len(ratios) == 1 else f"x{ratios[0]:.2f}-x{ratios[-1]:.2f}" scope_text = f"{compile_id} layer forward measured {span} compiled vs eager across {len(fwd_votes)} shape{'s' if len(fwd_votes) > 1 else ''}" kinds = {test for _ratio, test in fwd_votes} if "faster" in kinds and "not_faster" not in kinds: verdict_text = "keep it on" if compile_active in {True, None} else "[yellow]compiled is faster here, drop SDNQ_USE_TORCH_COMPILE=0[/yellow]" elif "not_faster" in kinds and "faster" not in kinds: verdict_text = "SDNQ_USE_TORCH_COMPILE=0 would help on this gpu" if compile_active in {True, None} else "keeping it off matches the measurement" elif "faster" in kinds: n_faster = sum(1 for _ratio, test in fwd_votes if test == "faster") n_slower = sum(1 for _ratio, test in fwd_votes if test == "not_faster") verdict_text = f"split verdict ({n_faster} faster, {n_slower} slower)" else: verdict_text = "within this run's noise" compile_note = f"dequant torch.compile ({compile_state}, {compile_origin}): {scope_text}; {verdict_text}" elif int8_eager and int8_compiled: compile_test = speed_verdict(int8_compiled, int8_eager, sigma=pair_sigma(entry("int8"), "compiled_ms", entry("int8"), "eager_ms"), margin=1.0) verdict_text = {"faster": "keep it on", "not_faster": "SDNQ_USE_TORCH_COMPILE=0 would help on this gpu", "inconclusive": "within this run's noise"}[compile_test] compile_note = f"dequant torch.compile ({compile_state}, {compile_origin}): int8 dequant kernel measured {ratio_text(int8_eager, int8_compiled)} compiled vs eager, no layer forward data; {verdict_text}" # judge quantized matmul on the dtype the current config would quantize with, falling back to int8; # speed never wins alone: the faster path must also hold output error within recommend_error_cap mm_entry = entry(mm_id) # explicit MatMul type rows are measured at the first dequant shape only; ignore them when # the recommendation reference fell back to another shape float_mm = float_mm_results if (float_mm_results and reference == dequant_shapes[0][0]) else {} def float_mm_entry(dtype_id, mm_dtype): return float_mm.get(f"{dtype_id}+{mm_dtype}") or {} # the enable checkbox is gone: one row on the type dropdown carries the verdict; # candidates are enabled (auto-resolved dtype) plus any measured explicit alternative mm_candidates = [] if mm_entry.get("mm_ms"): mm_candidates.append(("enabled", mm_entry.get("mm_dtype"), mm_entry["mm_ms"], mm_entry.get("mm_err"))) if mm_id in float_mm_dtypes: for alt_dtype in float_mm_alternative_dtypes: alt = float_mm_entry(mm_id, alt_dtype) if alt.get("mm_ms"): mm_candidates.append((alt_dtype, alt.get("mm_dtype"), alt["mm_ms"], alt.get("mm_err"))) if mm_candidates and mm_entry.get("fwd_ms"): fastest_mm = min(ms for _sel, _res, ms, _err in mm_candidates) near_fastest = [c for c in mm_candidates if c[2] <= fastest_mm * 1.05] best_sel, best_resolved, best_ms, best_err = min(near_fastest, key=lambda c: c[3] if c[3] is not None else float("inf")) fwd_err = mm_entry.get("fwd_err") err_ok = not (fwd_err and best_err) or best_err <= fwd_err * recommend_error_cap best_entry = mm_entry if best_sel == "enabled" else float_mm_entry(mm_id, best_sel) mm_test = speed_verdict(best_ms, mm_entry["fwd_ms"], sigma=pair_sigma(best_entry, "mm_ms", mm_entry, "fwd_ms"), z=sidak_z_for(len(mm_candidates))) recommend_mm = mm_test == "faster" and err_ok mm_reason = f"{mm_id} quantized matmul ({best_resolved}) measured {ratio_text(mm_entry['fwd_ms'], best_ms)} vs the dequant path" if fwd_err and best_err: mm_reason += f", output error {best_err:.5f} vs {fwd_err:.5f}" if mm_test == "faster" and not err_ok: mm_reason += f"; not recommended, the speedup multiplies output error beyond x{recommend_error_cap:.0f}" if weights_mode not in mode_ids: mm_reason += f"; current quantization type {weights_mode} was not benchmarked, judged on int8" type_parts = [] for sel, resolved, ms, err in mm_candidates: sel_label = f"enabled ({resolved})" if sel == "enabled" else sel type_parts.append(f"{sel_label} {ms:.3f} ms err {err:.5f}" if err is not None else f"{sel_label} {ms:.3f} ms") if len(mm_candidates) > 1: mm_reason += f"; lowest error within 5% of the fastest: {', '.join(type_parts)}" elif best_sel == "enabled": mm_reason += f"; enabled resolves to {best_resolved} for {mm_id} weights" else: mm_reason += f"; enabled (float8_e4m3fn) failed on this gpu for {mm_id} weights" accum_alt = float_mm.get(f"{mm_id}+float16+accum") or {} if accum_alt.get("mm_ms"): mm_reason += f"; float16 mm with SDNQ_TRITON_MM_USE_FP16_ACCUM=1 measured {accum_alt['mm_ms']:.3f} ms (env var, not a dropdown value)" if mm_test == "inconclusive" and err_ok: mm_choice = current("sdnq_quantize_matmul_mode") mm_reason += "; too close to the margin to call at this run's noise, keeping the current setting" else: mm_choice = best_sel if recommend_mm else "disabled" rows.append(("Quantized MatMul type", current("sdnq_quantize_matmul_mode"), mm_choice, mm_reason)) # the te override is independent now, with its own dtype option; judge it at # text-encoder geometry using the dtype the te config would quantize with te_shape = next((label for label, _out, _in in dequant_shapes if label.startswith("te ")), None) te_weights_mode = str(getattr(shared.opts, "sdnq_quantize_weights_mode_te", "Same as model")) if te_weights_mode in {"Same as model", "default"}: te_weights_mode = weights_mode te_mm_id = mode_ids.get(te_weights_mode, "int8") te_entry = ((dequant_results.get(te_shape) or {}).get(te_mm_id) or {}) if te_shape else {} if te_entry.get("fwd_ms") and te_entry.get("mm_ms"): te_fwd_err, te_mm_err = te_entry.get("fwd_err"), te_entry.get("mm_err") te_err_ok = not (te_fwd_err and te_mm_err) or te_mm_err <= te_fwd_err * recommend_error_cap te_test = speed_verdict(te_entry["mm_ms"], te_entry["fwd_ms"], sigma=pair_sigma(te_entry, "mm_ms", te_entry, "fwd_ms")) te_reason = f"{te_mm_id} quantized matmul measured {ratio_text(te_entry['fwd_ms'], te_entry['mm_ms'])} vs the dequant path at {te_shape}" if te_fwd_err and te_mm_err: te_reason += f", output error {te_mm_err:.5f} vs {te_fwd_err:.5f}" if te_test == "faster" and not te_err_ok: te_reason += f"; not recommended, the speedup multiplies output error beyond x{recommend_error_cap:.0f}" if te_test == "inconclusive" and te_err_ok: te_choice = current("sdnq_quantize_matmul_mode_te") te_reason += "; too close to the margin to call at this run's noise, keeping the current setting" else: te_choice = "enabled" if te_test == "faster" and te_err_ok else "disabled" te_reason += "; applies when text encoders are sdnq-quantized" rows.append(("Quantized MatMul type for Text Encoders", current("sdnq_quantize_matmul_mode_te"), te_choice, te_reason)) sweeps = sweep_results or {} sweeps_at_reference = reference == dequant_shapes[0][0] mm_on = getattr(shared.opts, "sdnq_quantize_matmul_mode", "disabled") != "disabled" # Group size: judged on the path the current config runs (mm cells when quantized matmul # is on); an explicit size must cut error meaningfully without real speed or size cost groups = sweeps.get("groups") or {} group_id = mm_id if mm_id in group_sweep_dtypes else "int8" auto_group = groups.get(f"{group_id}@0") or {} if sweeps_at_reference and auto_group: def group_metrics(g_entry): ms = g_entry.get("mm_ms") if mm_on else g_entry.get("fwd_ms") err = g_entry.get("mm_err") if mm_on else g_entry.get("weight_err") return ms, err auto_ms, auto_err = group_metrics(auto_group) candidates = [] for group in group_sweep_values: if group == 0: continue g_entry = groups.get(f"{group_id}@{group}") or {} g_ms, g_err = group_metrics(g_entry) if g_ms and g_err and g_entry.get("size_bytes"): candidates.append((group, g_ms, g_err, g_entry["size_bytes"])) if auto_ms and auto_err and auto_group.get("size_bytes"): qualifying = [c for c in candidates if c[2] < auto_err * 0.80 and c[1] <= auto_ms * 1.05 and c[3] <= auto_group["size_bytes"] * 1.10] resolved = auto_group.get("mm_resolved") if mm_on else auto_group.get("resolved") if qualifying: best_group, best_gms, best_gerr, best_gsize = min(qualifying, key=lambda c: c[2]) group_reason = f"{group_id} error {best_gerr:.5f} vs {auto_err:.5f} at auto ({resolved}), {best_gms:.3f} vs {auto_ms:.3f} ms, {size_cell(best_gsize).strip()} vs {size_cell(auto_group['size_bytes']).strip()}" rows.append(("Group size", current("sdnq_group_size"), str(best_group), group_reason)) else: explicit_best = min(candidates, key=lambda c: c[2]) if candidates else None group_reason = f"auto resolves to {resolved} for {group_id} weights" if mm_on and explicit_best and auto_err and abs(explicit_best[2] - auto_err) <= auto_err * 0.10: group_reason += f"; grouped storage is re-quantized per forward for quantized matmul, so no swept group changed mm error (best {explicit_best[2]:.5f} vs {auto_err:.5f})" elif explicit_best: group_reason += f"; best explicit ({explicit_best[0]}) measured error {explicit_best[2]:.5f} vs {auto_err:.5f} at {explicit_best[1]:.3f} vs {auto_ms:.3f} ms and {size_cell(explicit_best[3]).strip()} vs {size_cell(auto_group['size_bytes']).strip()}" rows.append(("Group size", current("sdnq_group_size"), "0", group_reason)) # Dequantize using full precision: off keeps scales in the model dtype; the speed gain # must hold error within the cap to be recommended toggles = sweeps.get("toggles") or {} fp32_off = toggles.get("fp32_off") or {} toggle_id = mm_id if mm_id in toggle_dtypes else "int8" off_entry = fp32_off.get(toggle_id) or {} on_entry = entry(toggle_id) if sweeps_at_reference and off_entry.get("fwd_ms") and on_entry.get("fwd_ms") and off_entry.get("weight_err") and on_entry.get("weight_err"): off_faster = off_entry["fwd_ms"] <= on_entry["fwd_ms"] * 0.95 off_err_ok = off_entry["weight_err"] <= on_entry["weight_err"] * recommend_error_cap recommend_fp32 = not (off_faster and off_err_ok) fp32_reason = f"{toggle_id} with full precision off: {off_entry['fwd_ms']:.3f} vs {on_entry['fwd_ms']:.3f} ms, weight error {off_entry['weight_err']:.5f} vs {on_entry['weight_err']:.5f}" rows.append(("Dequantize using full precision", current("sdnq_dequantize_fp32"), str(recommend_fp32), fp32_reason)) # conv settings: the conv matmul verdict is spatial-size dependent, so every measured conv # shape weighs in; the recommended value follows the first (unet-class) shape, with the # other shape's number in the reason so vae-quantizing sessions can differ conv_all = sweeps.get("conv") or {} conv_per_shape = {} for conv_shape_label, _oc, _ic, _k, _px in conv_shapes: shape_conv = conv_all.get(conv_shape_label) or {} mm_rows_here = [(config_id, shape_conv.get(config_id) or {}) for config_id in ("int8-mm", "uint8-mm")] conv_per_shape[conv_shape_label] = (shape_conv.get("bf16") or {}, shape_conv.get("int8") or {}, [(config_id, row) for config_id, row in mm_rows_here if row.get("fwd_ms")]) conv_base, conv_int8, conv_mm_rows = conv_per_shape.get(conv_shapes[0][0], ({}, {}, [])) if conv_base.get("fwd_ms"): quant_rows = [(config_id, row) for config_id, row in [("int8", conv_int8)] + conv_mm_rows if row.get("fwd_ms")] if quant_rows: best_conv_id, best_conv = min(quant_rows, key=lambda item: item[1]["fwd_ms"]) conv_free = best_conv["fwd_ms"] <= conv_base["fwd_ms"] * 1.05 conv_reason = f"{best_conv_id} conv measured {best_conv['fwd_ms']:.3f} vs {conv_base['fwd_ms']:.3f} ms {dtype_label()} at {conv_shapes[0][0]}, output error {best_conv.get('out_err', 0):.5f} vs {conv_base.get('out_err', 0):.5f}, size {size_cell(best_conv.get('size_bytes', 0)).strip()} vs {size_cell(conv_base.get('size_bytes', 0)).strip()}" if not conv_free: conv_reason += "; enable for the size saving, not speed" rows.append(("Quantize convolutional layers", current("sdnq_quantize_conv_layers"), str(conv_free), conv_reason)) if conv_int8.get("fwd_ms") and conv_mm_rows: best_mm_id, best_mm = min(conv_mm_rows, key=lambda item: item[1]["fwd_ms"]) conv_mm_test = speed_verdict(best_mm["fwd_ms"], conv_int8["fwd_ms"], sigma=pair_sigma(best_mm, "fwd_ms", conv_int8, "fwd_ms"), z=sidak_z_for(len(conv_mm_rows))) mm_err_ok = not (best_mm.get("out_err") and conv_int8.get("out_err")) or best_mm["out_err"] <= conv_int8["out_err"] * recommend_error_cap conv_mm_reason = f"{best_mm_id} measured {best_mm['fwd_ms']:.3f} vs {conv_int8['fwd_ms']:.3f} ms for the int8 conv dequant path at {conv_shapes[0][0]}, output error {best_mm.get('out_err', 0):.5f} vs {conv_int8.get('out_err', 0):.5f}" other_base, _other_int8, other_mm_rows = conv_per_shape.get(conv_shapes[1][0], ({}, {}, [])) if len(conv_shapes) > 1 else ({}, {}, []) if other_base.get("fwd_ms") and other_mm_rows: other_best_id, other_best = min(other_mm_rows, key=lambda item: item[1]["fwd_ms"]) conv_mm_reason += f"; at {conv_shapes[1][0]} {other_best_id} measured {ratio_text(other_base['fwd_ms'], other_best['fwd_ms'])} vs the {dtype_label()} conv" if other_best["fwd_ms"] > other_base["fwd_ms"] * 1.05: conv_mm_reason += ", a slowdown; disable when quantizing vae-class convs" if conv_mm_test == "inconclusive" and mm_err_ok: conv_mm_choice = current("sdnq_use_quantized_matmul_conv") conv_mm_reason += "; too close to the margin to call at this run's noise, keeping the current setting" else: conv_mm_choice = str(conv_mm_test == "faster" and mm_err_ok) rows.append(("Use quantized MatMul with conv", current("sdnq_use_quantized_matmul_conv"), conv_mm_choice, conv_mm_reason)) differing = 0 for setting, current_value, recommended, reason in rows: if current_value == recommended: table.add_row(setting, f"[green]{current_value}[/green]", recommended, reason) else: differing += 1 table.add_row(setting, f"[yellow]{current_value}[/yellow]", recommended, reason) emit(table) if differing: emit(f"[yellow]{differing} setting{'s' if differing > 1 else ''} differ{'' if differing > 1 else 's'} from the recommended value, change in Compute Settings -> SDNQ[/yellow]") elif rows: emit("[green]current settings already match the recommendations[/green]") notes = [] if compile_note: notes.append(compile_note) speedups = [] for dtype_id, label, _cfg in dequant_dtype_configs: text = ratio_text(entry(dtype_id).get("eager_ms"), entry(dtype_id).get("compiled_ms")) if text: speedups.append(f"{label} {text}") if speedups: notes.append(f"compiled dequant speedup vs eager at {reference}: {', '.join(speedups)}; standalone kernel scope, the compile verdict uses the layer forward") fp8_entry, fp8sdnq_entry = entry("fp8"), entry("fp8sdnq") if fp8_entry.get("compiled_ms") and fp8sdnq_entry.get("compiled_ms"): notes.append(f"fp8 storage: native e4m3 compiled {ratio_text(fp8sdnq_entry['compiled_ms'], fp8_entry['compiled_ms'])} vs the float8_e4m3fn_sdnq uint8 view") elif fp8sdnq_entry.get("compiled_ms") and fp8_entry.get("eager_ms"): notes.append(f"fp8 storage: float8_e4m3fn_sdnq (uint8 view) compiled {ratio_text(fp8_entry['eager_ms'], fp8sdnq_entry['compiled_ms'])} vs eager e4m3 dequant") if weight_dequant_result is not None: e5m2_ok = weight_dequant_result["float8_e5m2"][0] e4m3_ok = weight_dequant_result["float8_e4m3fn"][0] if not e4m3_ok and e5m2_ok: notes.append("compiled dequant: e5m2 compiles raw, e4m3 does not; sdnq's upcast covers e4m3") elif not e4m3_ok and not e5m2_ok: notes.append("[yellow]compiled dequant: neither fp8 dtype compiles raw; sdnq upcasts e4m3, e5m2 has no such path[/yellow]") fp8_mm_failed = [dtype_id for dtype_id, _label, _cfg in dequant_dtype_configs if entry(dtype_id).get("mm_error") and entry(dtype_id).get("mm_dtype") == "float8_e4m3fn"] if fp8_mm_failed: notes.append(f"[yellow]enabled quantized matmul auto-selects float8_e4m3fn for {', '.join(fp8_mm_failed)} and failed; set an explicit MatMul type for float weight dtypes[/yellow]") for dtype_id in fp8_mm_failed: alt_parts = [] for alt_dtype in float_mm_alternative_dtypes: alt = float_mm_entry(dtype_id, alt_dtype) if alt.get("mm_ms"): part = f"{alt_dtype} mm {alt['mm_ms']:.3f} ms" if alt.get("mm_err") is not None: part += f" err {alt['mm_err']:.5f}" alt_parts.append(part) dequant_fwd, dequant_err = entry(dtype_id).get("fwd_ms"), entry(dtype_id).get("fwd_err") if alt_parts and dequant_fwd: base_part = f"dequant path {dequant_fwd:.3f} ms" if dequant_err is not None: base_part += f" err {dequant_err:.5f}" notes.append(f"{dtype_id} with MatMul type set explicitly: {', '.join(alt_parts)}; {base_part}") if base_ms and entry(mm_id).get("fwd_ms"): notes.append(f"current quantization type {weights_mode}: dequant-path forward {ratio_text(entry(mm_id)['fwd_ms'], base_ms)} vs a {dtype_label()} nn.Linear at {reference}") # size/error frontier: the dtypes no other measured dtype beats on both axes at once sized = [(dtype_id, e["size_bytes"], e["weight_err"]) for dtype_id, _label, _cfg in dequant_dtype_configs for e in [entry(dtype_id)] if e.get("size_bytes") and e.get("weight_err")] frontier = sorted((c for c in sized if not any(o is not c and o[1] <= c[1] and o[2] <= c[2] for o in sized)), key=lambda c: c[1]) if len(frontier) > 1: notes.append(f"size/error frontier at {reference}: " + ", ".join(f"{dtype_id} {size_cell(size).strip()} err {err:.3f}" for dtype_id, size, err in frontier)) if variant_results: for variant_id, _cfg in dequant_variant_configs: parts = [] for dtype_id in dequant_variant_dtypes: variant_entry = variant_results.get(f"{dtype_id}+{variant_id}") or {} if variant_entry.get("err_ratio"): parts.append(f"{dtype_id} err x{variant_entry['err_ratio']:.2f} for {quant_cell(variant_entry['quant_s']).strip()} quant") if parts: notes.append(f"{variant_id} vs plain quantization: {', '.join(parts)}") groups_sweep = sweeps.get("groups") or {} if groups_sweep and sweeps_at_reference: group_parts = [] mm_flat = True for dtype_id in group_sweep_dtypes: auto_g = groups_sweep.get(f"{dtype_id}@0") or {} swept = [(group, groups_sweep.get(f"{dtype_id}@{group}") or {}) for group in group_sweep_values if group != 0] swept = [(group, g_entry) for group, g_entry in swept if g_entry.get("weight_err")] if auto_g.get("weight_err") and swept: best_group, best_entry = min(swept, key=lambda item: item[1]["weight_err"]) if best_entry.get("size_bytes") and auto_g.get("size_bytes"): group_parts.append(f"{dtype_id} group {best_group} weight err x{best_entry['weight_err'] / auto_g['weight_err']:.2f} of auto for {size_cell(best_entry['size_bytes']).strip()} vs {size_cell(auto_g['size_bytes']).strip()}") mm_errs = [g_entry["mm_err"] for _group, g_entry in swept if g_entry.get("mm_err")] + ([auto_g["mm_err"]] if auto_g.get("mm_err") else []) if len(mm_errs) > 1 and max(mm_errs) > min(mm_errs) * 1.10: mm_flat = False if group_parts: groups_note = "group size storage tradeoff (dequant path): " + ", ".join(group_parts) if mm_flat: groups_note += "; grouped storage is re-quantized per forward for quantized matmul, no swept group changed mm error" notes.append(groups_note) conv_note_all = sweeps.get("conv") or {} conv_note_parts = [] for conv_shape_label, _oc, _ic, _k, _px in conv_shapes: shape_conv = conv_note_all.get(conv_shape_label) or {} base_row = shape_conv.get("bf16") or {} mm_here = [row for config_id in ("int8-mm", "uint8-mm") for row in [shape_conv.get(config_id) or {}] if row.get("fwd_ms")] if base_row.get("fwd_ms") and mm_here: best_here = min(mm_here, key=lambda row: row["fwd_ms"]) conv_note_parts.append(f"{conv_shape_label} best conv mm {ratio_text(base_row['fwd_ms'], best_here['fwd_ms'])}") if len(conv_note_parts) > 1: notes.append(f"conv quantized matmul vs the {dtype_label()} conv by shape: " + "; ".join(conv_note_parts)) te_shape = next((label for label, _out, _in in dequant_shapes if label.startswith("te ")), None) te_results = dequant_results.get(te_shape) or {} if te_shape else {} te_int8 = te_results.get("int8") or {} if te_int8.get("fwd_ms") and te_int8.get("mm_ms") and entry("int8").get("fwd_ms") and entry("int8").get("mm_ms"): notes.append(f"text-encoder geometry ({te_shape}): int8 quantized mm {ratio_text(te_int8['fwd_ms'], te_int8['mm_ms'])} vs the dequant path, {ratio_text(entry('int8')['fwd_ms'], entry('int8')['mm_ms'])} at {reference}") krea2_parts = [] for krea2_label, _out, _in in dequant_shapes: if not krea2_label.startswith("krea2 "): continue krea2_int8 = (dequant_results.get(krea2_label) or {}).get("int8") or {} if krea2_int8.get("fwd_ms") and krea2_int8.get("mm_ms"): krea2_parts.append(f"{krea2_label} {ratio_text(krea2_int8['fwd_ms'], krea2_int8['mm_ms'])}") if krea2_parts and entry("int8").get("fwd_ms") and entry("int8").get("mm_ms"): notes.append(f"krea 2 geometry: int8 quantized mm vs the dequant path {', '.join(krea2_parts)}; {ratio_text(entry('int8')['fwd_ms'], entry('int8')['mm_ms'])} at {reference}") svd_sweep = sweeps.get("svd") or {} if svd_sweep and sweeps_at_reference: for dtype_id in svd_rank_sweep_dtypes: parts = [] for rank in svd_rank_sweep_values: rank_entry = svd_sweep.get(f"{dtype_id}@{rank}") or {} if rank_entry.get("err_ratio"): parts.append(f"rank {rank} err x{rank_entry['err_ratio']:.2f} ({size_cell(rank_entry['size_bytes']).strip()})") if parts: notes.append(f"svd rank on {dtype_id}: {', '.join(parts)}") hgroup_sweep = sweeps.get("hgroups") or {} hgroup_parts = [f"group {hgroup} err x{hentry['err_ratio']:.2f}" for hgroup in hadamard_group_values for hentry in [hgroup_sweep.get(str(hgroup)) or {}] if hentry.get("err_ratio")] if hgroup_parts and sweeps_at_reference: notes.append(f"weight-side hadamard group size on int8: {', '.join(hgroup_parts)}") dynamic_sweep = (sweeps.get("toggles") or {}).get("dynamic") or {} dynamic_parts = [] for requested in dynamic_quant_requests: dyn_entry = dynamic_sweep.get(requested) or {} if dyn_entry.get("chosen"): dynamic_parts.append(f"{requested} -> {dyn_entry['chosen']}" + (f" err {dyn_entry['weight_err']:.3f}" if dyn_entry.get("weight_err") else "")) if dynamic_parts: notes.append(f"dynamic quantization on these synthetic weights (default loss thresholds): {', '.join(dynamic_parts)}") cpu_quant_s = (sweeps.get("toggles") or {}).get("cpu_quant_s") gpu_quant_s = (sweeps.get("toggles") or {}).get("gpu_quant_s") if cpu_quant_s and gpu_quant_s: notes.append(f"Quantize using GPU: one int8 layer at {dequant_shapes[0][0]} takes {quant_cell(gpu_quant_s).strip()} on gpu vs {quant_cell(cpu_quant_s).strip()} on cpu; scales with layer count at load") notes.append("weight err: relative error of the dequantized weight vs fp32; forward outputs inherit it") notes.append("errors use synthetic outlier-heavy weights; rotation and svd gains are outlier-driven, real models gain less") if notes: emit(Panel("\n".join(notes), title="dequant notes", box=ROUNDED_BOX)) report["recommendations_dequant"] = dict(reference=reference, rows=[list(row) for row in rows], notes=notes) def save_report(path, args, sections, selected): run_info = dict(timestamp=time.strftime("%Y-%m-%dT%H:%M:%S"), sections=sections, shapes=selected, iters=args.iters, warmup=args.warmup, dtype=dtype_label()) run_info.update(report.get("run") or {}) # keep the drift fields main writes after the bench sections report["run"] = run_info with open(path, "w", encoding="utf-8") as fh: json.dump(report, fh, indent=2, default=str) console.print(f"structured results saved to {path}") pending_outputs = None # set by main once arguments are validated, so a crash can still flush outputs_written = False def flush_outputs(): # write --save/--json exactly once, from the normal path or the crash handler; partial # results beat losing an hour-long run to a kernel fault in one config global outputs_written # pylint: disable=global-statement if outputs_written or pending_outputs is None: return outputs_written = True if pending_outputs["save"]: save_transcript(pending_outputs["save"]) if pending_outputs["json"]: save_report(pending_outputs["json"], pending_outputs["args"], pending_outputs["sections"], pending_outputs["selected"]) def main(): global bench_dtype # pylint: disable=global-statement args = parse_cli() sections = [s.strip().lower() for s in args.sections.split(",") if s.strip()] unknown_sections = [s for s in sections if s not in all_sections] if unknown_sections: console.print(f"[red]unknown section(s): {', '.join(unknown_sections)}; available: {', '.join(all_sections)}[/red]") sys.exit(1) known_dtypes = [dtype_id for dtype_id, _label, _cfg in dequant_dtype_configs] selected_dtypes = None if args.dequant_dtypes.strip().lower() == "all" else [s.strip() for s in args.dequant_dtypes.split(",") if s.strip()] if selected_dtypes is not None: unknown_dtypes = [s for s in selected_dtypes if s not in known_dtypes] if unknown_dtypes: console.print(f"[red]unknown dequant dtype(s): {', '.join(unknown_dtypes)}; available: {', '.join(known_dtypes)}[/red]") sys.exit(1) known_blocks = [config_id for config_id, _w, _mm, _a in block_configs] selected_blocks = None if args.block_configs.strip().lower() == "all" else [s.strip() for s in args.block_configs.split(",") if s.strip()] selected_geometries = None if args.block_geometries.strip().lower() == "all" else [s.strip() for s in args.block_geometries.split(",") if s.strip()] if selected_geometries is not None: unknown_geometries = [s for s in selected_geometries if s not in block_geometries] if unknown_geometries: console.print(f"[red]unknown block geometry(ies): {', '.join(unknown_geometries)}; available: {', '.join(block_geometries)}[/red]") sys.exit(1) if selected_blocks is not None: unknown_blocks = [s for s in selected_blocks if s not in known_blocks] if unknown_blocks: console.print(f"[red]unknown block config(s): {', '.join(unknown_blocks)}; available: {', '.join(known_blocks)}[/red]") sys.exit(1) known_attention = [config_id for config_id, _label, _kwargs in bench_configs] selected_attention = None if args.configs.strip().lower() == "all" else [s.strip() for s in args.configs.split(",") if s.strip()] if selected_attention is not None: unknown_attention = [s for s in selected_attention if s not in known_attention] if unknown_attention: console.print(f"[red]unknown attention config(s): {', '.join(unknown_attention)}; available: {', '.join(known_attention)}[/red]") sys.exit(1) known_variants = [variant_id for variant_id, _cfg in dequant_variant_configs] variants_arg = args.dequant_variants.strip().lower() if variants_arg == "all": selected_variants = list(known_variants) elif variants_arg == "none": selected_variants = [] else: selected_variants = [s.strip() for s in args.dequant_variants.split(",") if s.strip()] unknown_variants = [s for s in selected_variants if s not in known_variants] if unknown_variants: console.print(f"[red]unknown dequant variant(s): {', '.join(unknown_variants)}; available: {', '.join(known_variants)}, all, none[/red]") sys.exit(1) sweeps_arg = args.dequant_sweeps.strip().lower() if sweeps_arg == "all": selected_sweeps = list(all_dequant_sweeps) elif sweeps_arg == "none": selected_sweeps = [] else: selected_sweeps = [s.strip() for s in args.dequant_sweeps.split(",") if s.strip()] unknown_sweeps = [s for s in selected_sweeps if s not in all_dequant_sweeps] if unknown_sweeps: console.print(f"[red]unknown dequant sweep(s): {', '.join(unknown_sweeps)}; available: {', '.join(all_dequant_sweeps)}, all, none[/red]") sys.exit(1) backends_arg = args.mm_backends.strip().lower() if backends_arg in {"none", ""}: selected_mm_backends = [] elif backends_arg == "all": selected_mm_backends = list(all_mm_backends) else: selected_mm_backends = [s.strip() for s in args.mm_backends.split(",") if s.strip()] unknown_backends = [s for s in selected_mm_backends if s not in all_mm_backends] if unknown_backends: console.print(f"[red]unknown matmul backend(s): {', '.join(unknown_backends)}; available: {', '.join(all_mm_backends)}, all, none[/red]") sys.exit(1) if not load_sdnext(): sys.exit(1) if args.dtype == "auto": # webui generation dtype: bf16 where native, fp16 on older gpus where bf16 is emulated if isinstance(getattr(devices, "dtype", None), torch.dtype) and devices.dtype in (torch.bfloat16, torch.float16): bench_dtype = devices.dtype else: bench_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16}[args.dtype] shapes_arg = args.shapes.strip().lower() shape_run_aliases = {"all": full_run, "sparse": sparse_run, "gate": gate_run} selected = list(shape_run_aliases[shapes_arg]) if shapes_arg in shape_run_aliases else [s.strip() for s in args.shapes.split(",") if s.strip()] unknown = [s for s in selected if s not in shape_presets] if unknown: console.print(f"[red]unknown shape preset(s): {', '.join(unknown)}; available: {', '.join(shape_presets)}[/red]") sys.exit(1) resolve_output_paths(args) global pending_outputs # pylint: disable=global-statement pending_outputs = dict(save=args.save, json=args.json, args=args, sections=sections, selected=selected) print_banner(selected, sections, args) prep_status, prep_detail = probe_compiled_prep() # must run first: see the cold-start note in probe_compiled_prep fp8_result = probe_fp8() if "attention" in sections else None weight_dequant_result = probe_weight_dequant_compile() if "dequant" in sections else None print_environment(fp8_result, prep_status, prep_detail, weight_dequant_result=weight_dequant_result) if not args.skip_checks and "attention" in sections: run_correctness() if report.get("correctness_aborted_after") and not cuda_context_alive(): flush_outputs() sys.exit(1) if args.skip_bench: flush_outputs() return dequant_results = {} if "dequant" in sections: if free_vram_gb() < 2.0: emit(f"[yellow]skipping dequant benchmarks: needs about 2 gb free vram, {free_vram_gb():.1f} gb available[/yellow]") else: from sdnq.common import check_torch_compile if not check_torch_compile(): # the module-level compiled dequant is a passthrough with the option off; swap in # a real compiled variant so the compiled fwd rows measure what enabling it gives from sdnq import dequantizer as dequantizer_module dequantizer_module.dequantize_weight_compiled = get_compiled_dequantize_weight() emit("[yellow]sdnq torch.compile is off (SDNQ_USE_TORCH_COMPILE=0): compiled fwd rows are measured with a tool-compiled dequant, matching the webui with it back on[/yellow]") for index, (shape_label, out_features, in_features) in enumerate(dequant_shapes, start=1): dequant_results[shape_label] = bench_dequant_shape(shape_label, out_features, in_features, args.iters, args.warmup, position=(index, len(dequant_shapes)), config_timeout=args.config_timeout, selected_dtypes=selected_dtypes) torch_device_module.empty_cache() variant_results = {} if selected_variants: first_label, first_out, first_in = dequant_shapes[0] variant_results = bench_dequant_variants(first_label, first_out, first_in, dequant_results.get(first_label, {}), selected_dtypes, selected_variants, args.iters, args.warmup, config_timeout=args.config_timeout) first_label, first_out, first_in = dequant_shapes[0] first_results = dequant_results.get(first_label, {}) float_mm_results = bench_float_mm_alternatives(first_label, first_out, first_in, first_results, selected_dtypes, args.iters, args.warmup, config_timeout=args.config_timeout) if selected_mm_backends: bench_mm_backends(first_label, first_out, first_in, selected_dtypes, selected_mm_backends, args.iters, args.warmup, config_timeout=args.config_timeout, rounds=args.mm_rounds) sweep_results = {} if "groups" in selected_sweeps: sweep_results["groups"] = bench_group_sizes(first_label, first_out, first_in, first_results, selected_dtypes, args.iters, args.warmup, config_timeout=args.config_timeout) if "svd" in selected_sweeps: sweep_results["svd"] = bench_svd_ranks(first_label, first_out, first_in, first_results, selected_dtypes, args.iters, args.warmup, config_timeout=args.config_timeout) if "hgroups" in selected_sweeps: sweep_results["hgroups"] = bench_hadamard_groups(first_label, first_out, first_in, first_results, selected_dtypes, args.iters, args.warmup, config_timeout=args.config_timeout) if "toggles" in selected_sweeps: sweep_results["toggles"] = bench_quant_toggles(first_label, first_out, first_in, first_results, selected_dtypes, args.iters, args.warmup, config_timeout=args.config_timeout) if "conv" in selected_sweeps: sweep_results["conv"] = bench_conv_section(args.iters, args.warmup, config_timeout=args.config_timeout) build_dequant_recommendations(dequant_results, weight_dequant_result, variant_results, float_mm_results, sweep_results) # the block section runs before the attention section: a broken-compile attention # environment disables dynamo globally, which would poison the block compiles if "block" in sections: if free_vram_gb() < 3.0: emit(f"[yellow]skipping block benchmarks: needs about 3 gb free vram, {free_vram_gb():.1f} gb available[/yellow]") else: bench_block_section(args.iters, args.warmup, config_timeout=args.timeout_flag, selected=selected_blocks, geometries=selected_geometries) if "attention" not in sections: emit_block_splits() if "attention" in sections: # bench the prep mode the advice points to: compiled, static workaround, or eager if prep_status == "failing_dynamic": from sdnq.kernels import triton_atten as atten_module inner = getattr(atten_module.get_attn_inputs, "_torchdynamo_orig_callable", None) atten_module.get_attn_inputs = torch.compile(inner, fullgraph=True, dynamic=False) emit("[yellow]dynamic-shape compile is broken here: benchmarking with the dynamic=false workaround applied, numbers match the webui after setting SDNQ_COMPILE_KWARGS='{\"dynamic\": false}'[/yellow]") elif prep_status == "failing": emit("[yellow]torch compile is broken here: benchmarking with eager input prep, numbers match the webui after setting SDNQ_USE_TORCH_COMPILE=0[/yellow]") from sdnq.kernels import triton_atten as atten_module inner = getattr(atten_module.get_attn_inputs, "_torchdynamo_orig_callable", None) if inner is not None: # swap in the eager prep; a disable toggle raises on torch 2.13+ atten_module.get_attn_inputs = inner all_results = {} minimum_vram = {"wan22": 6.0, "wan22-cfg": 12.0, "ltx2": 3.0, "masked": 6.0, "krea2": 8.0} for index, preset in enumerate(selected, start=1): needed = minimum_vram.get(preset, 2.0) if free_vram_gb() < needed: emit(f"[yellow]skipping {preset}: needs about {needed:.0f} gb free vram, {free_vram_gb():.1f} gb available[/yellow]") continue all_results[preset] = bench_shape(preset, args.iters, args.warmup, position=(index, len(selected)), config_timeout=args.timeout_flag, fp8_result=fp8_result, selected=selected_attention) emit_block_splits() build_recommendations(all_results, fp8_result, prep_status, block_results=(report.get("block") or {}).get("results"), block_variants=report.get("blocks")) if drift_samples: report.setdefault("run", {})["drift_sigma"] = run_drift_sigma() report["run"]["drift_samples"] = len(drift_samples) emit(f"[dim]run drift: sentinel re-measurements moved {run_drift_sigma():.1%} rms across {len(drift_samples)} shape windows; verdicts fold this into their uncertainty[/dim]") flush_outputs() if __name__ == "__main__": try: main() except KeyboardInterrupt: console.print("[yellow]interrupted, saving partial results[/yellow]") flush_outputs() sys.exit(130) except Exception as main_error: # pylint: disable=broad-exception-caught console.print(f"[red]benchmark aborted: {error_summary(main_error, 200)}[/red]") if not cuda_context_alive(): console.print("[red]the cuda context is corrupted (kernel faults are sticky); partial results saved, a rerun is required to continue[/red]") flush_outputs() raise