From ea95d50660cbe0d2b4ccf8c63134b3cfb5668251 Mon Sep 17 00:00:00 2001 From: papertager <2567587994@qq.com> Date: Fri, 5 Jun 2026 15:52:31 +0800 Subject: [PATCH 1/6] Add FlashMLA GPU smoke runner --- README.md | 13 ++++++ tools/run_flash_mla_smoke.py | 77 ++++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+) create mode 100644 tools/run_flash_mla_smoke.py diff --git a/README.md b/README.md index 976379be..4f5bca12 100644 --- a/README.md +++ b/README.md @@ -39,6 +39,19 @@ python setup.py install python tests/test_flash_mla.py ``` +### Smoke test + +After building the extension, run a small correctness 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, then runs a +single bf16 FlashMLA case against the PyTorch reference implementation. Use +`--dtype fp16` or the shape flags in `--help` to cover additional cases. + ### Usage ```python diff --git a/tools/run_flash_mla_smoke.py b/tools/run_flash_mla_smoke.py new file mode 100644 index 00000000..85312924 --- /dev/null +++ b/tools/run_flash_mla_smoke.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +import argparse +import random +from pathlib import Path +import sys + +import torch + + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from tests.test_flash_mla import test_flash_mla # 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 correctness smoke test on one MACA device." + ) + 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") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + device = torch.device(args.device) + if device.type != "cuda": + raise ValueError("FlashMLA smoke test requires a CUDA-compatible MACA device.") + + 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}") + + test_flash_mla( + b=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, + causal=not args.non_causal, + varlen=args.varlen, + block_size=args.block_size, + ) + print("flash_mla_smoke_ok") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From ea23c721d25cea2f820c356aab4963fc5419d1b6 Mon Sep 17 00:00:00 2001 From: papertager <2567587994@qq.com> Date: Mon, 22 Jun 2026 13:12:09 +0800 Subject: [PATCH 2/6] fix: validate FlashMLA smoke device early --- tools/run_flash_mla_smoke.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tools/run_flash_mla_smoke.py b/tools/run_flash_mla_smoke.py index 85312924..3617ea89 100644 --- a/tools/run_flash_mla_smoke.py +++ b/tools/run_flash_mla_smoke.py @@ -43,9 +43,18 @@ def parse_args() -> argparse.Namespace: 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) From 6f1912624868a718e21960b335864cc49a09895d Mon Sep 17 00:00:00 2001 From: ghangz <152254226+ghangz@users.noreply.github.com> Date: Mon, 29 Jun 2026 17:21:20 +0800 Subject: [PATCH 3/6] sync review update for README.md --- README.md | 183 +++++++++++++++++++++++++++--------------------------- 1 file changed, 93 insertions(+), 90 deletions(-) diff --git a/README.md b/README.md index 4f5bca12..5ff53385 100644 --- a/README.md +++ b/README.md @@ -1,91 +1,94 @@ -# FlashMLA on MXMACA -We provide the implementation of FlashMLA from FlashAttention-2(version 2.6.3), based on MACA toolkit and C500 chips. - -FlashAttention-2 currently supports: -1. Datatype fp16 and bf16. -2. Multi-Token Prediction greater or equal to 1. -3. Paged kvcache with block size equal to 2^n (n >= 0) - -## How to run on MXMACA Device -## Installation - -Requirements: -- MXMACA GPUs. -- MACA development toolkit. -- mcTlass source code. -- mcPytorch2.1 and mcTriton2.1 from maca toolkit wheel package and above. - -To install flash attn in conda env: -1. Make sure that maca pytorch2.1 and triton2.1 is installed. -2. Download mctlass source code from FlashAttention2/csrc/mctlass on website : https://sw-download.metax-tech.com/ - -### Set environment variables -```bash -export MACA_PATH=/your/maca/path -export CUDA_PATH=$MACA_PATH/tools/cu-bridge -export MACA_CLANG_PATH=$MACA_PATH/mxgpu_llvm/bin -export LD_LIBRARY_PATH=$MACA_PATH/lib:$MACA_PATH/mxgpu_llvm/lib:$MACA_PATH/ompi/lib:$LD_LIBRARY_PATH -``` - -### Install - -```bash -python setup.py install -``` - -### Benchmark - -```bash -python tests/test_flash_mla.py -``` - -### Smoke test - -After building the extension, run a small correctness 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, then runs a -single bf16 FlashMLA case against the PyTorch reference implementation. Use -`--dtype fp16` or the shape flags in `--help` to cover additional cases. - -### Usage - -```python -from flash_mla import get_mla_metadata, flash_mla_with_kvcache - -tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv) - -for i in range(num_layers): - ... - o_i, lse_i = flash_mla_with_kvcache( - q_i, kvcache_i, block_table, cache_seqlens, dv, - tile_scheduler_metadata, num_splits, causal=True, - ) - ... -``` - -## Requirements - -- MXMACA 2.27 and above -- PyTorch 2.0 and above - -## Acknowledgement - -FlashMLA is inspired by [FlashAttention 2&3](https://github.com/dao-AILab/flash-attention/) and [cutlass](https://github.com/nvidia/cutlass) projects. - - - -## Citation - -```bibtex -@misc{flashmla2025, - title={FlashMLA: Efficient MLA decoding kernel}, - author={Jiashi Li}, - year={2025}, - publisher = {GitHub}, - howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}}, +# FlashMLA on MXMACA +We provide the implementation of FlashMLA from FlashAttention-2(version 2.6.3), based on MACA toolkit and C500 chips. + +FlashAttention-2 currently supports: +1. Datatype fp16 and bf16. +2. Multi-Token Prediction greater or equal to 1. +3. Paged kvcache with block size equal to 2^n (n >= 0) + +## How to run on MXMACA Device +## Installation + +Requirements: +- MXMACA GPUs. +- MACA development toolkit. +- mcTlass source code. +- mcPytorch2.1 and mcTriton2.1 from maca toolkit wheel package and above. + +To install flash attn in conda env: +1. Make sure that maca pytorch2.1 and triton2.1 is installed. +2. Download mctlass source code from FlashAttention2/csrc/mctlass on website : https://sw-download.metax-tech.com/ + +### Set environment variables +```bash +export MACA_PATH=/your/maca/path +export CUDA_PATH=$MACA_PATH/tools/cu-bridge +export MACA_CLANG_PATH=$MACA_PATH/mxgpu_llvm/bin +export LD_LIBRARY_PATH=$MACA_PATH/lib:$MACA_PATH/mxgpu_llvm/lib:$MACA_PATH/ompi/lib:$LD_LIBRARY_PATH +``` + +### Install + +```bash +python setup.py install +``` + +### Benchmark + +```bash +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 +from flash_mla import get_mla_metadata, flash_mla_with_kvcache + +tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv) + +for i in range(num_layers): + ... + o_i, lse_i = flash_mla_with_kvcache( + q_i, kvcache_i, block_table, cache_seqlens, dv, + tile_scheduler_metadata, num_splits, causal=True, + ) + ... +``` + +## Requirements + +- MXMACA 2.27 and above +- PyTorch 2.0 and above + +## Acknowledgement + +FlashMLA is inspired by [FlashAttention 2&3](https://github.com/dao-AILab/flash-attention/) and [cutlass](https://github.com/nvidia/cutlass) projects. + + + +## Citation + +```bibtex +@misc{flashmla2025, + title={FlashMLA: Efficient MLA decoding kernel}, + author={Jiashi Li}, + year={2025}, + publisher = {GitHub}, + howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}}, } From 3429398a10b6b7a36f1b5c52254e1bc88141a346 Mon Sep 17 00:00:00 2001 From: ghangz <152254226+ghangz@users.noreply.github.com> Date: Mon, 29 Jun 2026 17:21:22 +0800 Subject: [PATCH 4/6] sync review update for tools/run_flash_mla_smoke.py --- tools/run_flash_mla_smoke.py | 283 ++++++++++++++++++++++++----------- 1 file changed, 198 insertions(+), 85 deletions(-) diff --git a/tools/run_flash_mla_smoke.py b/tools/run_flash_mla_smoke.py index 3617ea89..bd9eac12 100644 --- a/tools/run_flash_mla_smoke.py +++ b/tools/run_flash_mla_smoke.py @@ -1,86 +1,199 @@ -#!/usr/bin/env python3 -import argparse -import random -from pathlib import Path -import sys - -import torch - - -REPO_ROOT = Path(__file__).resolve().parents[1] -if str(REPO_ROOT) not in sys.path: - sys.path.insert(0, str(REPO_ROOT)) - -from tests.test_flash_mla import test_flash_mla # 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 correctness smoke test on one MACA device." - ) - 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") - return parser.parse_args() - - -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}") - - test_flash_mla( - b=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, - causal=not args.non_causal, - varlen=args.varlen, - block_size=args.block_size, - ) - print("flash_mla_smoke_ok") - return 0 - - -if __name__ == "__main__": +#!/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()) From b5253d35333a10ed0c102177e343deb245c46199 Mon Sep 17 00:00:00 2001 From: ghangz <152254226+ghangz@users.noreply.github.com> Date: Mon, 29 Jun 2026 17:22:07 +0800 Subject: [PATCH 5/6] sync review update: README.md --- README.md | 186 +++++++++++++++++++++++++++--------------------------- 1 file changed, 93 insertions(+), 93 deletions(-) diff --git a/README.md b/README.md index 5ff53385..5cdef98d 100644 --- a/README.md +++ b/README.md @@ -1,94 +1,94 @@ -# FlashMLA on MXMACA -We provide the implementation of FlashMLA from FlashAttention-2(version 2.6.3), based on MACA toolkit and C500 chips. - -FlashAttention-2 currently supports: -1. Datatype fp16 and bf16. -2. Multi-Token Prediction greater or equal to 1. -3. Paged kvcache with block size equal to 2^n (n >= 0) - -## How to run on MXMACA Device -## Installation - -Requirements: -- MXMACA GPUs. -- MACA development toolkit. -- mcTlass source code. -- mcPytorch2.1 and mcTriton2.1 from maca toolkit wheel package and above. - -To install flash attn in conda env: -1. Make sure that maca pytorch2.1 and triton2.1 is installed. -2. Download mctlass source code from FlashAttention2/csrc/mctlass on website : https://sw-download.metax-tech.com/ - -### Set environment variables -```bash -export MACA_PATH=/your/maca/path -export CUDA_PATH=$MACA_PATH/tools/cu-bridge -export MACA_CLANG_PATH=$MACA_PATH/mxgpu_llvm/bin -export LD_LIBRARY_PATH=$MACA_PATH/lib:$MACA_PATH/mxgpu_llvm/lib:$MACA_PATH/ompi/lib:$LD_LIBRARY_PATH -``` - -### Install - -```bash -python setup.py install -``` - -### Benchmark - -```bash -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 -from flash_mla import get_mla_metadata, flash_mla_with_kvcache - -tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv) - -for i in range(num_layers): - ... - o_i, lse_i = flash_mla_with_kvcache( - q_i, kvcache_i, block_table, cache_seqlens, dv, - tile_scheduler_metadata, num_splits, causal=True, - ) - ... -``` - -## Requirements - -- MXMACA 2.27 and above -- PyTorch 2.0 and above - -## Acknowledgement - -FlashMLA is inspired by [FlashAttention 2&3](https://github.com/dao-AILab/flash-attention/) and [cutlass](https://github.com/nvidia/cutlass) projects. - - - -## Citation - -```bibtex -@misc{flashmla2025, - title={FlashMLA: Efficient MLA decoding kernel}, - author={Jiashi Li}, - year={2025}, - publisher = {GitHub}, - howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}}, +# FlashMLA on MXMACA +We provide the implementation of FlashMLA from FlashAttention-2(version 2.6.3), based on MACA toolkit and C500 chips. + +FlashAttention-2 currently supports: +1. Datatype fp16 and bf16. +2. Multi-Token Prediction greater or equal to 1. +3. Paged kvcache with block size equal to 2^n (n >= 0) + +## How to run on MXMACA Device +## Installation + +Requirements: +- MXMACA GPUs. +- MACA development toolkit. +- mcTlass source code. +- mcPytorch2.1 and mcTriton2.1 from maca toolkit wheel package and above. + +To install flash attn in conda env: +1. Make sure that maca pytorch2.1 and triton2.1 is installed. +2. Download mctlass source code from FlashAttention2/csrc/mctlass on website : https://sw-download.metax-tech.com/ + +### Set environment variables +```bash +export MACA_PATH=/your/maca/path +export CUDA_PATH=$MACA_PATH/tools/cu-bridge +export MACA_CLANG_PATH=$MACA_PATH/mxgpu_llvm/bin +export LD_LIBRARY_PATH=$MACA_PATH/lib:$MACA_PATH/mxgpu_llvm/lib:$MACA_PATH/ompi/lib:$LD_LIBRARY_PATH +``` + +### Install + +```bash +python setup.py install +``` + +### Benchmark + +```bash +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 +from flash_mla import get_mla_metadata, flash_mla_with_kvcache + +tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv) + +for i in range(num_layers): + ... + o_i, lse_i = flash_mla_with_kvcache( + q_i, kvcache_i, block_table, cache_seqlens, dv, + tile_scheduler_metadata, num_splits, causal=True, + ) + ... +``` + +## Requirements + +- MXMACA 2.27 and above +- PyTorch 2.0 and above + +## Acknowledgement + +FlashMLA is inspired by [FlashAttention 2&3](https://github.com/dao-AILab/flash-attention/) and [cutlass](https://github.com/nvidia/cutlass) projects. + + + +## Citation + +```bibtex +@misc{flashmla2025, + title={FlashMLA: Efficient MLA decoding kernel}, + author={Jiashi Li}, + year={2025}, + publisher = {GitHub}, + howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}}, } From 030e1ad516f51cf2e861f417f1a5f23b8e6e642b Mon Sep 17 00:00:00 2001 From: ghangz <152254226+ghangz@users.noreply.github.com> Date: Mon, 29 Jun 2026 17:22:09 +0800 Subject: [PATCH 6/6] sync review update: tools/run_flash_mla_smoke.py --- tools/run_flash_mla_smoke.py | 396 +++++++++++++++++------------------ 1 file changed, 198 insertions(+), 198 deletions(-) diff --git a/tools/run_flash_mla_smoke.py b/tools/run_flash_mla_smoke.py index bd9eac12..3285f00a 100644 --- a/tools/run_flash_mla_smoke.py +++ b/tools/run_flash_mla_smoke.py @@ -1,199 +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__": +#!/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())