diff --git a/README.md b/README.md index 976379be..5cdef98d 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,22 @@ python setup.py install python tests/test_flash_mla.py ``` +### Smoke test + +After building the extension, run a small smoke case before launching the full +benchmark suite: + +```bash +python tools/run_flash_mla_smoke.py +``` + +The command prints the detected torch version and MACA device name, executes one +FlashMLA case, and reports the output and LSE deltas against the PyTorch +reference implementation. This is useful for validating that the extension can +compile, launch, and return numerically stable values on a target MACA stack. +Use `--dtype fp16`, the shape flags in `--help`, or optional thresholds such as +`--max-lse-cos-diff 1e-4` to tailor the check for your environment. + ### Usage ```python diff --git a/tools/run_flash_mla_smoke.py b/tools/run_flash_mla_smoke.py new file mode 100644 index 00000000..3285f00a --- /dev/null +++ b/tools/run_flash_mla_smoke.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +import argparse +import json +import math +import random +from pathlib import Path +import sys + +import torch +import triton + + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from flash_mla import flash_mla_with_kvcache, get_mla_metadata # noqa: E402 +from tests.test_flash_mla import scaled_dot_product_attention # noqa: E402 + + +def _dtype(value: str) -> torch.dtype: + if value == "bf16": + return torch.bfloat16 + if value == "fp16": + return torch.float16 + raise argparse.ArgumentTypeError("dtype must be bf16 or fp16") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run a small FlashMLA smoke test and report reference deltas." + ) + parser.add_argument("--device", default="cuda:0", help="Torch device to run on.") + parser.add_argument("--dtype", type=_dtype, default=torch.bfloat16, help="bf16 or fp16.") + parser.add_argument("--batch-size", type=int, default=128) + parser.add_argument("--s-q", type=int, default=1) + parser.add_argument("--mean-sk", type=int, default=4096) + parser.add_argument("--h-q", type=int, default=16) + parser.add_argument("--h-kv", type=int, default=1) + parser.add_argument("--d", type=int, default=576) + parser.add_argument("--dv", type=int, default=512) + parser.add_argument("--block-size", type=int, default=16) + parser.add_argument("--varlen", action="store_true") + parser.add_argument("--non-causal", action="store_true") + parser.add_argument("--json", action="store_true", help="Print the final summary as JSON.") + parser.add_argument("--max-out-cos-diff", type=float, default=None) + parser.add_argument("--max-lse-cos-diff", type=float, default=None) + return parser.parse_args() + + +def _metric_dict(x: torch.Tensor, y: torch.Tensor) -> dict[str, float]: + x = x.double() + y = y.double() + rmse = ((x - y) * (x - y)).mean().sqrt().item() + cos_diff = 1 - 2 * (x * y).sum().item() / max((x * x + y * y).sum().item(), 1e-12) + amax_diff = (x - y).abs().max().item() + return { + "rmse": rmse, + "cos_diff": cos_diff, + "amax_diff": amax_diff, + } + + +def _run_flash_mla_case(args: argparse.Namespace) -> tuple[dict[str, object], torch.Tensor, torch.Tensor]: + cache_seqlens = torch.full((args.batch_size,), args.mean_sk, dtype=torch.int32) + if args.varlen: + for i in range(args.batch_size): + cache_seqlens[i] = max(random.normalvariate(args.mean_sk, args.mean_sk / 2), args.s_q) + max_seqlen = cache_seqlens.max().item() + max_seqlen_pad = triton.cdiv(max_seqlen, 256) * 256 + + q = torch.randn(args.batch_size, args.s_q, args.h_q, args.d) + block_table = torch.arange( + args.batch_size * max_seqlen_pad // args.block_size, dtype=torch.int32 + ).view(args.batch_size, max_seqlen_pad // args.block_size) + blocked_k = torch.randn(block_table.numel(), args.block_size, args.h_kv, args.d) + for i in range(args.batch_size): + blocked_k.view(args.batch_size, max_seqlen_pad, args.h_kv, args.d)[ + i, cache_seqlens[i].item() : + ] = float("nan") + blocked_v = blocked_k[..., : args.dv] + + tile_scheduler_metadata, num_splits = get_mla_metadata( + cache_seqlens, args.s_q * args.h_q // args.h_kv, args.h_kv + ) + out_flash, lse_flash = flash_mla_with_kvcache( + q, + blocked_k, + block_table, + cache_seqlens, + args.dv, + tile_scheduler_metadata, + num_splits, + causal=not args.non_causal, + ) + + out_ref = torch.empty(args.batch_size, args.s_q, args.h_q, args.dv, dtype=torch.float32) + lse_ref = torch.empty(args.batch_size, args.h_q, args.s_q, dtype=torch.float32) + blocked_k_view = blocked_k.view(-1, args.h_kv, args.d) + blocked_v_view = blocked_v.view(-1, args.h_kv, args.dv) + for i in range(args.batch_size): + begin = i * max_seqlen_pad + end = begin + cache_seqlens[i] + out_i, lse_i = scaled_dot_product_attention( + q[i].transpose(0, 1), + blocked_k_view[begin:end].transpose(0, 1), + blocked_v_view[begin:end].transpose(0, 1), + h_q=args.h_q, + h_kv=args.h_kv, + is_causal=not args.non_causal, + ) + out_ref[i] = out_i.transpose(0, 1) + lse_ref[i] = lse_i + + summary = { + "shape": { + "batch_size": args.batch_size, + "s_q": args.s_q, + "mean_sk": args.mean_sk, + "h_q": args.h_q, + "h_kv": args.h_kv, + "d": args.d, + "dv": args.dv, + "block_size": args.block_size, + "varlen": args.varlen, + "causal": not args.non_causal, + }, + "cache_seqlens": { + "min": int(cache_seqlens.min().item()), + "max": int(cache_seqlens.max().item()), + "mean": float(cache_seqlens.float().mean().item()), + }, + "out": _metric_dict(out_flash, out_ref), + "lse": _metric_dict(lse_flash, lse_ref), + } + return summary, out_flash, lse_flash + + +def _enforce_thresholds(args: argparse.Namespace, summary: dict[str, object]) -> None: + checks = [] + if args.max_out_cos_diff is not None: + checks.append(("out.cos_diff", summary["out"]["cos_diff"], args.max_out_cos_diff)) + if args.max_lse_cos_diff is not None: + checks.append(("lse.cos_diff", summary["lse"]["cos_diff"], args.max_lse_cos_diff)) + failed = [(name, value, limit) for name, value, limit in checks if value > limit] + if failed: + details = ", ".join(f"{name}={value:.6g} > {limit:.6g}" for name, value, limit in failed) + raise SystemExit(f"flash_mla_smoke_threshold_failed: {details}") + + +def main() -> int: + args = parse_args() + if not torch.cuda.is_available(): + raise RuntimeError( + "CUDA is not available. Please check your MACA driver and PyTorch installation." + ) + device = torch.device(args.device) + if device.type != "cuda": + raise ValueError("FlashMLA smoke test requires a CUDA-compatible MACA device.") + if device.index is not None and device.index >= torch.cuda.device_count(): + raise ValueError( + f"Device index {device.index} is out of range. Total available devices: " + f"{torch.cuda.device_count()}" + ) + + torch.set_default_dtype(args.dtype) + torch.set_default_device(device) + torch.cuda.set_device(device) + torch.manual_seed(0) + random.seed(0) + + print(f"torch={torch.__version__}") + print(f"device={torch.cuda.get_device_name(device)}") + print(f"dtype={args.dtype}") + summary, out_flash, lse_flash = _run_flash_mla_case(args) + print( + "out_metrics=" + f"cos_diff={summary['out']['cos_diff']:.6g}, " + f"rmse={summary['out']['rmse']:.6g}, " + f"amax_diff={summary['out']['amax_diff']:.6g}" + ) + print( + "lse_metrics=" + f"cos_diff={summary['lse']['cos_diff']:.6g}, " + f"rmse={summary['lse']['rmse']:.6g}, " + f"amax_diff={summary['lse']['amax_diff']:.6g}" + ) + print(f"flash_out_shape={tuple(out_flash.shape)}") + print(f"flash_lse_shape={tuple(lse_flash.shape)}") + if args.json: + print(json.dumps(summary, indent=2, sort_keys=True)) + _enforce_thresholds(args, summary) + + print("flash_mla_smoke_ok") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())