This package contains both the trimmed Python/Triton reproducer and a CUDA Driver API + PTX variant that does not require Python, PyTorch, Triton, or Inductor at runtime.
Both variants require an SM100 GPU and a compatible CUDA driver. They were verified with bundled ptxas-blackwell 13.3 (V13.3.33) and driver 580.126.20 on GB200.
Run repro.py with an affected Triton/ptxas installation:
This doc servers as a quick reference for the _scaled_mm API and how it has changed overtime for each major version of PyTorch.
NOTE The leading underscore is intended here and we make no current FC/BC guarantees on this API. That being said it is currently the only OP that has native support for FP8 matmuls within the PyTorch Libary. We are planning to make an official Public api for this. Until then this is subject to change but you can use this doc as a reference.
| import itertools | |
| from collections import defaultdict | |
| from contextlib import nullcontext | |
| from dataclasses import asdict, dataclass | |
| from typing import Callable, List, Tuple | |
| from tabulate import tabulate | |
| from tqdm import tqdm | |
| import torch |
| import torch | |
| torch.set_float32_matmul_precision("high") | |
| import torch.utils.benchmark as benchmark | |
| from diffusers import DiffusionPipeline | |
| import gc | |
| # from torchao.quantization import ( | |
| # int4_weight_only, |
| #!/usr/bin/env bash | |
| # Install newest clang with `bash -c "$(wget -O - https://apt.llvm.org/llvm.sh)"` | |
| # chmod u+x update_alternatives_clang.sh | |
| # ./update_alternatives_clang.sh <version> <priority> | |
| update_alternatives() { | |
| local version=${1} | |
| local priority=${2} | |
| local master=${3} | |
| local slaves=${4} |