#!/usr/bin/env python3 """Correctness tests for fp8-gemm.""" from __future__ import annotations import argparse import importlib import json import math import os import sys from dataclasses import asdict, dataclass from pathlib import Path import torch ROOT = Path(__file__).resolve().parents[2] PACKAGE = ROOT / "fp8-gemm" REGISTRATION_INCLUDE = ( ROOT.parent / "kernels" / "kernel-builder" / "src" / "pyproject" / "templates" / "torch" ) SHAPES = { "decode_m1_k512_n512": (1, 512, 512), "decode_m1_k4096_n2048": (1, 4096, 2048), "decode_m1_k4096_n8192": (1, 4096, 8192), "small_m8_k1024_n2048": (8, 1024, 2048), "small_m16_k4096_n4096": (16, 4096, 4096), "small_m32_k4096_n8192": (32, 4096, 8192), "small_m64_k512_n1024": (64, 512, 1024), } SM110_SHAPES = { "large_m_boundary_65": (65, 2048, 2048), # PI0.5 / PI0 decoder and encoder projection families. "pi05_action_qkv": (51, 2048, 2560), "pi05_action_o": (51, 2048, 2048), "pi05_action_gate_up": (51, 2048, 16384), "pi05_action_down": (51, 8192, 2048), # GROOT N1.6/N1.7 DiT, backbone, and vision rows. "groot_dit_qkv": (51, 1536, 4608), "groot_n17_llm_o": (277, 2048, 2048), "groot_n17_llm_gate_up": (277, 2048, 16384), "groot_n17_llm_down": (277, 8192, 2048), "groot_n17_vit_o": (1024, 1024, 1024), # Cosmos Edge and LingBot projection families. "cosmos_edge_action": (64, 2048, 9216), "lingbot_vision_o": (1024, 1280, 1280), "lingbot_action_gate_up": (105, 2048, 16384), # PI0.5 Thor prefill tower, full real row envelope. "pi05_prefill_qkv": (712, 2048, 2560), "pi05_prefill_o": (970, 2048, 2048), "pi05_prefill_gate_up": (768, 2048, 32768), "pi05_prefill_down": (768, 16384, 2048), } MODES = { "smoke": ["decode_m1_k512_n512", "small_m8_k1024_n2048"], "headline": [ "decode_m1_k4096_n2048", "decode_m1_k4096_n8192", "small_m16_k4096_n4096", "small_m32_k4096_n8192", ], "full": list(SHAPES.keys()), } @dataclass class Metrics: shape: str M: int K: int N: int variant: int tile: str max_abs: float mean_abs: float p99_abs: float cosine: float dtype: str tolerance: str passed: bool class SourceOps: def __init__(self, namespace: str) -> None: self._ops = getattr(torch.ops, namespace) @staticmethod def select_fp8_linear_tile(m: int, n: int, k: int, variant: int = 0) -> str: return select_tile(m, n, k, variant) def fp8_linear_bf16(self, x, w, alpha=1.0, out=None, variant=0): if out is None: out = torch.empty((x.shape[0], w.shape[0]), device=x.device, dtype=torch.bfloat16) self._ops.fp8_linear_bf16(x, w, float(alpha), int(variant), out) return out def fp8_linear_residual_bf16(self, x, w, residual, alpha=1.0, variant=0): self._ops.fp8_linear_residual_bf16(x, w, float(alpha), int(variant), residual) return residual def fp8_linear_bias_bf16(self, x, w, bias, alpha=1.0, out=None): if out is None: out = torch.empty( (x.shape[0], w.shape[0]), device=x.device, dtype=torch.bfloat16 ) self._ops.fp8_linear_bias_bf16(x, w, bias, float(alpha), out) return out def fp8_linear_bias_residual_bf16( self, x, w, bias, residual, alpha=1.0 ): self._ops.fp8_linear_bias_residual_bf16( x, w, bias, float(alpha), residual ) return residual def fp8_linear_bias_gelu_bf16(self, x, w, bias, alpha=1.0, out=None): if out is None: out = torch.empty( (x.shape[0], w.shape[0]), device=x.device, dtype=torch.bfloat16 ) self._ops.fp8_linear_bias_gelu_bf16(x, w, bias, float(alpha), out) return out def fp8_blockwise_linear_bf16( self, x, w, input_scale, weight_scale, out=None ): if out is None: out = torch.empty( (x.shape[0], w.shape[0]), device=x.device, dtype=torch.bfloat16, ) self._ops.fp8_blockwise_linear_bf16( x, w, input_scale, weight_scale, out ) return out def fp8_blockwise_swiglu_quantize_fp8( self, x, gate_up_weight, input_scale, gate_up_weight_scale, output=None, output_scale=None, ): n = gate_up_weight.shape[0] // 2 if output is None: output = torch.empty( (x.shape[0], n), device=x.device, dtype=torch.float8_e4m3fn ) if output_scale is None: output_scale = torch.empty( (x.shape[0], n // 128), device=x.device, dtype=torch.float32 ) self._ops.fp8_blockwise_swiglu_quantize_fp8( x, gate_up_weight, input_scale, gate_up_weight_scale, output, output_scale, ) return output, output_scale def _current_arch_list() -> str: major, minor = torch.cuda.get_device_capability(0) if (major, minor) == (11, 0): return "11.0a" return "12.0a" if (major, minor) == (12, 0) else f"{major}.{minor}" def load_source_ops() -> SourceOps: from torch.utils.cpp_extension import load if not REGISTRATION_INCLUDE.is_dir(): raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}") os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _current_arch_list()) namespace = "fp8_gemm_source_test" cutlass_include = Path( os.environ.get( "CUTLASS_INCLUDE", str( ROOT.parent / "flashrt_pr31_review" / "third_party" / "cutlass" / "include" ), ) ) if not (cutlass_include / "cutlass" / "cutlass.h").is_file(): raise RuntimeError( "CUTLASS 4 include path is required; set CUTLASS_INCLUDE" ) capability = torch.cuda.get_device_capability(0) if capability == (8, 9): cuda_sources = [ str(PACKAGE / "csrc" / "fp8_block128_gemm_mma_sm89.cu"), str(PACKAGE / "csrc" / "fp8_gemv_m1_sm89.cu"), ] source_define = "-DFLASHRT_FP8_GEMM_SOURCE_SM89_ONLY" elif capability == (11, 0): cuda_sources = [ str(PACKAGE / "csrc" / "cutlass_sm110_fp8_gemm.cu"), str(PACKAGE / "csrc" / "cublaslt_fp8_bias_sm110.cu"), ] source_define = "-DFLASHRT_FP8_GEMM_SOURCE_SM110_ONLY" else: cuda_sources = [ str(PACKAGE / "csrc" / "fp8_gemv_m1_sm120.cu"), str(PACKAGE / "csrc" / "fp8_smallM_handtuned_sm120.cu"), str(PACKAGE / "csrc" / "fp8_smallM_handtuned_ldmatrix_sm120.cu"), str(PACKAGE / "csrc" / "cutlass_sm120_block128_fp8_gemm.cu"), ] source_define = "-DFLASHRT_FP8_GEMM_SOURCE_SM120_ONLY" load( name=namespace, sources=[str(PACKAGE / "torch-ext" / "torch_binding.cpp"), *cuda_sources], extra_include_paths=[ str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE), str(cutlass_include), str(cutlass_include.parent / "tools" / "util" / "include"), ], extra_cflags=["-O3", "-DNDEBUG", "-DCUDA_KERNEL", source_define], extra_cuda_cflags=[ "-O3", "-DNDEBUG", "--expt-relaxed-constexpr", "--use_fast_math", "-U__CUDA_NO_HALF_OPERATORS__", "-U__CUDA_NO_HALF_CONVERSIONS__", "-U__CUDA_NO_BFLOAT16_CONVERSIONS__", "-U__CUDA_NO_HALF2_OPERATORS__", "-DCUDA_KERNEL", source_define ], verbose=False, ) return SourceOps(namespace) def load_installed_ops(artifact: str | None): if artifact: sys.path.insert(0, artifact) try: return importlib.import_module("fp8_gemm") finally: if artifact: sys.path.remove(artifact) def select_tile(m: int, n: int, k: int, variant: int = 0) -> str: if torch.cuda.get_device_capability(0) == (11, 0): forced = {1: "sm110_sq_bf16", 2: "sm110_t1_bf16", 3: "sm110_wide_bf16"} if variant not in {0, *forced}: raise RuntimeError("SM110 variant must be in [0, 3]") if variant: return forced[variant] if m >= 512 and k == 2048 and 2048 <= n <= 2560: return "sm110_sq_bf16" if m >= 512 and n >= 16 * k: return "sm110_t1_bf16" if m >= 512 and k >= 4 * n: return "sm110_wide_bf16" if n >= 8 * k: return "sm110_wide_bf16" if m >= 128 and k >= 4 * n: return "sm110_sq_bf16" if n == k and m >= 512: return "sm110_sq_bf16" if k <= 1024 else "sm110_wide_bf16" if n == k and m >= 128: return "sm110_wide_bf16" return "sm110_t1_bf16" if m == 1: if variant == 4: return "gemv_fp8_m1_w4" if variant == 8: return "gemv_fp8_m1_w8" if variant == 16: return "gemv_fp8_m1_w16" if n <= 2048: return "gemv_fp8_m1_w4" if n <= 8192: return "gemv_fp8_m1_w8" return "gemv_fp8_m1_w16" if m <= 16: if k % 256 == 0: return "ld_fp8_gemm_16x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_16x64x256_w4" if n % 256 == 0: return "ld_fp8_gemm_16x256x128_w8" if n % 192 == 0: return "ld_fp8_gemm_16x192x128_w4" if n % 128 == 0: return "ld_fp8_gemm_16x128x128_w4" return "ld_fp8_gemm_16x64x128_w4" if m <= 32: if k % 256 == 0: return "ld_fp8_gemm_32x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_32x64x256_w4" if n % 192 == 0: return "ld_fp8_gemm_32x192x128_w4" if n % 128 == 0: return "ld_fp8_gemm_32x128x128_w4" return "ld_fp8_gemm_32x64x128_w4" if m <= 64: if k % 256 == 0: return "ld_fp8_gemm_64x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_64x64x256_w4" if n % 128 == 0: return "ld_fp8_gemm_64x128x128_w4" return "ld_fp8_gemm_64x64x128_w4" if m <= 64: if k % 256 == 0: return "ld_fp8_gemm_64x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_64x64x256_w4" if n % 128 == 0: return "ld_fp8_gemm_64x128x128_w4" return "ld_fp8_gemm_64x64x128_w4" raise RuntimeError("unsupported M") def make_inputs(m: int, k: int, n: int, seed: int): gen = torch.Generator(device="cuda") gen.manual_seed(seed) x_bf16 = (torch.randn((m, k), device="cuda", generator=gen) * 0.25).to(torch.bfloat16) w_bf16 = (torch.randn((n, k), device="cuda", generator=gen) * 0.25).to(torch.bfloat16) x = x_bf16.to(torch.float8_e4m3fn) w = w_bf16.to(torch.float8_e4m3fn) return x, w def reference(x: torch.Tensor, w: torch.Tensor, alpha: float) -> torch.Tensor: return ((x.float() @ w.float().T) * float(alpha)).to(torch.bfloat16) def compare(got: torch.Tensor, expected: torch.Tensor) -> tuple[float, float, float, float]: diff = (got.float() - expected.float()).abs().flatten() max_abs = float(diff.max().item()) mean_abs = float(diff.mean().item()) p99_rank = max(1, min(diff.numel(), math.ceil(0.99 * diff.numel()))) p99_abs = float(diff.kthvalue(p99_rank).values.item()) cos = float(torch.nn.functional.cosine_similarity(got.float().flatten(), expected.float().flatten(), dim=0).item()) return max_abs, mean_abs, p99_abs, cos def check_threshold(max_abs: float, mean_abs: float, p99_abs: float, cos: float) -> bool: return max_abs <= 0.5 and mean_abs <= 0.02 and p99_abs <= 0.25 and cos >= 0.999 def run_case(ops, name: str, shape: tuple[int, int, int], variant: int = 0) -> Metrics: m, k, n = shape x, w = make_inputs(m, k, n, seed=1000 + m + k + n + variant) alpha = 1.0 expected = reference(x, w, alpha) got = ops.fp8_linear_bf16(x, w, alpha=alpha, variant=variant) torch.cuda.synchronize() max_abs, mean_abs, p99_abs, cos = compare(got, expected) tile = ops.select_fp8_linear_tile(m, n, k, variant) passed = check_threshold(max_abs, mean_abs, p99_abs, cos) return Metrics( shape=name, M=m, K=k, N=n, variant=variant, tile=tile, max_abs=max_abs, mean_abs=mean_abs, p99_abs=p99_abs, cosine=cos, dtype=str(got.dtype), tolerance="max_abs<=0.5 mean_abs<=0.02 p99_abs<=0.25 cosine>=0.999", passed=passed, ) def run_residual_case(ops) -> Metrics: m, k, n = (1, 4096, 2048) x, w = make_inputs(m, k, n, seed=2026) residual = torch.randn((1, n), device="cuda", dtype=torch.bfloat16) * 0.1 expected = (residual.float() + reference(x, w, 1.0).float()).to(torch.bfloat16) got = residual.clone() variant = 0 if torch.cuda.get_device_capability(0) == (11, 0) else 8 ops.fp8_linear_residual_bf16(x, w, got, alpha=1.0, variant=variant) torch.cuda.synchronize() max_abs, mean_abs, p99_abs, cos = compare(got, expected) passed = check_threshold(max_abs, mean_abs, p99_abs, cos) return Metrics( shape="decode_residual_m1_k4096_n2048", M=m, K=k, N=n, variant=variant, tile=( "sm110_t1_bf16_residual" if torch.cuda.get_device_capability(0) == (11, 0) else "gemv_fp8_m1_resadd_w8" ), max_abs=max_abs, mean_abs=mean_abs, p99_abs=p99_abs, cosine=cos, dtype=str(got.dtype), tolerance="max_abs<=0.5 mean_abs<=0.02 p99_abs<=0.25 cosine>=0.999", passed=passed, ) def run_bias_cases(ops) -> int: count = 0 shapes = [ (512, 1152, 4304), (768, 4304, 1152), (768, 1152, 3456), ] for m, k, n in shapes: x, w = make_inputs(m, k, n, seed=7000 + m + k + n) bias = (torch.randn((n,), device="cuda") * 0.1).to(torch.bfloat16) alpha = 0.75 base = (x.float() @ w.float().T) * alpha got = ops.fp8_linear_bias_bf16(x, w, bias, alpha=alpha) expected = (base + bias.float()).to(torch.bfloat16) maximum, mean, p99, cosine = compare(got, expected) assert maximum <= 0.5 and mean <= 0.02 and p99 <= 0.25 and cosine >= 0.999, ( "bias", m, k, n, maximum, mean, p99, cosine ) residual = (torch.randn((m, n), device="cuda") * 0.1).to( torch.bfloat16 ) residual_before = residual.clone() got_residual = ops.fp8_linear_bias_residual_bf16( x, w, bias, residual, alpha=alpha ) expected_residual = ( residual_before.float() + base + bias.float() ).to(torch.bfloat16) maximum, mean, p99, cosine = compare(got_residual, expected_residual) assert maximum <= 0.5 and mean <= 0.02 and p99 <= 0.25 and cosine >= 0.999, ( "bias_residual", m, k, n, maximum, mean, p99, cosine ) got_gelu = ops.fp8_linear_bias_gelu_bf16(x, w, bias, alpha=alpha) expected_gelu = torch.nn.functional.gelu( base + bias.float(), approximate="tanh" ).to(torch.bfloat16) maximum, mean, p99, cosine = compare(got_gelu, expected_gelu) assert maximum <= 0.5 and mean <= 0.02 and p99 <= 0.25 and cosine >= 0.999, ( "bias_gelu", m, k, n, maximum, mean, p99, cosine ) count += 3 m, k, n = (512, 1152, 4304) x, w = make_inputs(m, k, n, seed=8801) bias = torch.randn((n,), device="cuda", dtype=torch.bfloat16) def invoke(input, weight, bias): return ops.fp8_linear_bias_bf16(input, weight, bias) eager = invoke(x, w, bias) compiled = torch.compile(invoke, fullgraph=True)(x, w, bias) torch.testing.assert_close(compiled, eager, rtol=0.0, atol=0.0) graph_out = torch.empty_like(eager) ops.fp8_linear_bias_bf16(x, w, bias, out=graph_out) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): ops.fp8_linear_bias_bf16(x, w, bias, out=graph_out) graph.replay() torch.testing.assert_close(graph_out, eager, rtol=0.0, atol=0.0) return count + 2 def run_blockwise_case( ops, name: str, shape: tuple[int, int, int] ) -> Metrics: m, k, n = shape gen = torch.Generator(device="cuda").manual_seed(5000 + m + k + n) x = (torch.randn((m, k), device="cuda", generator=gen) * 0.4).to( torch.float8_e4m3fn ) w = (torch.randn((n, k), device="cuda", generator=gen) * 0.4).to( torch.float8_e4m3fn ) input_scale = ( 0.005 + 0.02 * torch.rand((m, k // 128), device="cuda", generator=gen) ).float().contiguous() weight_scale = ( 0.005 + 0.02 * torch.rand((n // 128, k // 128), device="cuda", generator=gen) ).float().contiguous() expanded_input_scale = input_scale.repeat_interleave(128, dim=1) expanded_weight_scale = weight_scale.repeat_interleave( 128, dim=0 ).repeat_interleave(128, dim=1) expected = ( (x.float() * expanded_input_scale) @ (w.float() * expanded_weight_scale).T ).to(torch.bfloat16) got = ops.fp8_blockwise_linear_bf16( x, w, input_scale, weight_scale ) torch.cuda.synchronize() max_abs, mean_abs, p99_abs, cos = compare(got, expected) passed = ( max_abs <= 0.0625 and mean_abs <= 0.003 and p99_abs <= 0.015625 and cos >= 0.9999 ) return Metrics( shape=name, M=m, K=k, N=n, variant=0, tile=( "mma_sm89_block128" if torch.cuda.get_device_capability(0) == (8, 9) else "cutlass_sm120_block128" ), max_abs=max_abs, mean_abs=mean_abs, p99_abs=p99_abs, cosine=cos, dtype=str(got.dtype), tolerance=( "max_abs<=0.0625 mean_abs<=0.003 " "p99_abs<=0.015625 cosine>=0.9999" ), passed=passed, ) def run_blockwise_compile_case(ops) -> None: m, k, n = (51, 1536, 1536) gen = torch.Generator(device="cuda").manual_seed(9153) x = (torch.randn((m, k), device="cuda", generator=gen) * 0.4).to( torch.float8_e4m3fn ) w = (torch.randn((n, k), device="cuda", generator=gen) * 0.4).to( torch.float8_e4m3fn ) input_scale = torch.rand( (m, k // 128), device="cuda", generator=gen, dtype=torch.float32 ).mul_(0.02).add_(0.005) weight_scale = torch.rand( (n // 128, k // 128), device="cuda", generator=gen, dtype=torch.float32, ).mul_(0.02).add_(0.005) def invoke(input, weight, input_scale, weight_scale): return ops.fp8_blockwise_linear_bf16( input, weight, input_scale, weight_scale ) eager = invoke(x, w, input_scale, weight_scale) compiled = torch.compile(invoke, fullgraph=True)( x, w, input_scale, weight_scale ) torch.testing.assert_close(compiled, eager, rtol=0.0, atol=0.0) def run_sm89_swiglu_case(ops, m: int, n: int, k: int) -> None: gen = torch.Generator(device="cuda").manual_seed(8900 + m + n + k) x = (torch.randn((m, k), device="cuda", generator=gen) * 0.3).to( torch.float8_e4m3fn ) weight = ( torch.randn((2 * n, k), device="cuda", generator=gen) * 0.3 ).to(torch.float8_e4m3fn) input_scale = torch.rand( (m, k // 128), device="cuda", generator=gen ).mul_(0.02).add_(0.005) weight_scale = torch.rand( (2 * n // 128, k // 128), device="cuda", generator=gen ).mul_(0.02).add_(0.005) output, output_scale = ops.fp8_blockwise_swiglu_quantize_fp8( x, weight, input_scale, weight_scale ) expanded_x_scale = input_scale.repeat_interleave(128, dim=1) expanded_w_scale = weight_scale.repeat_interleave(128, dim=0).repeat_interleave(128, dim=1) x_f32 = x.float() * expanded_x_scale weight_f32 = weight.float() * expanded_w_scale gate, up = (x_f32 @ weight_f32.t()).split(n, dim=1) expected = ( torch.nn.functional.silu(gate).bfloat16() * up.bfloat16() ).bfloat16() actual = ( output.float() * output_scale.repeat_interleave(128, dim=1) ).bfloat16() maximum, mean, p99, cosine = compare(actual, expected) assert output.dtype == torch.float8_e4m3fn assert output_scale.dtype == torch.float32 assert torch.isfinite(output_scale).all() and (output_scale > 0).all() assert cosine >= 0.999 and mean <= 0.01 and p99 <= 0.05, ( m, n, k, maximum, mean, p99, cosine ) def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--backend", choices=["source", "installed"], default="source") parser.add_argument("--artifact", default=None) parser.add_argument("--mode", choices=sorted(MODES), default="smoke") parser.add_argument("--json-out", default=None) args = parser.parse_args() if not torch.cuda.is_available(): raise SystemExit("CUDA is required") capability = torch.cuda.get_device_capability(0) if capability not in {(8, 9), (11, 0), (12, 0)}: raise SystemExit( "fp8-gemm source tests require SM89, SM110, or SM120; " f"got SM{capability[0]}{capability[1]}" ) ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact) rows = [] if capability in {(11, 0), (12, 0)}: rows.extend(run_case(ops, name, SHAPES[name]) for name in MODES[args.mode]) rows.append(run_residual_case(ops)) if capability == (11, 0) and args.mode == "full": rows.extend( run_case(ops, name, shape) for name, shape in SM110_SHAPES.items() ) rows.extend( run_case( ops, f"sm110_forced_variant_{variant}", SM110_SHAPES["pi05_action_gate_up"], variant, ) for variant in (1, 2, 3) ) bias_count = run_bias_cases(ops) else: bias_count = 0 if capability in {(8, 9), (12, 0)}: blockwise_shapes = [ ("blockwise_decode", (1, 1024, 1024)), ("blockwise_action", (51, 1536, 1536)), ] if args.mode == "full": blockwise_shapes += [ ("blockwise_groot", (277, 2048, 2048)), ("blockwise_vision", (1024, 1152, 1152)), ("blockwise_video", (2520, 3072, 3072)), ("blockwise_qwen_mlp", (128, 4096, 12288)), ] rows.extend( run_blockwise_case(ops, name, shape) for name, shape in blockwise_shapes ) run_blockwise_compile_case(ops) if capability == (8, 9): for m, n, k in [ (1, 128, 128), (16, 512, 1024), (31, 1536, 1536), (32, 2048, 4096), (51, 4096, 4096), (128, 4096, 4096), (256, 4096, 4096), ]: run_sm89_swiglu_case(ops, m, n, k) try: x = torch.zeros((257, 128), device="cuda", dtype=torch.float8_e4m3fn) w = torch.zeros((256, 128), device="cuda", dtype=torch.float8_e4m3fn) xs = torch.ones((257, 1), device="cuda", dtype=torch.float32) ws = torch.ones((2, 1), device="cuda", dtype=torch.float32) ops.fp8_blockwise_swiglu_quantize_fp8(x, w, xs, ws) except RuntimeError as error: assert "M <= 256" in str(error) else: raise AssertionError("M=257 must be rejected") failed = [row for row in rows if not row.passed] payload = { "passed": len(rows) - len(failed) + bias_count, "failed": len(failed), "rows": [asdict(row) for row in rows], "bias_checks": bias_count, } print(json.dumps(payload, indent=2, sort_keys=True)) if args.json_out: output_path = Path(args.json_out) output_path.parent.mkdir(parents=True, exist_ok=True) output_path.write_text( json.dumps(payload, indent=2, sort_keys=True) + "\n" ) if failed: raise SystemExit(1) if __name__ == "__main__": main()