diff --git a/BENCHMARK_H200.md b/BENCHMARK_H200.md index 5c82b66..ab9676c 100644 --- a/BENCHMARK_H200.md +++ b/BENCHMARK_H200.md @@ -1,59 +1,118 @@ # Benchmark Results — Hopper (SM90) -> Auto-generated by `benchmarks/generate_benchmark_hopper_md.py` on 2026-05-19. +> Measured on 2026-08-09. -> **GPU:** NVIDIA H200 | **CUDA:** 12.9 | **PyTorch:** 2.9.1+cu129 +> **GPU:** NVIDIA H200 141GB, 132 SMs | **CUDA:** 12.9 | **PyTorch:** 2.9.1+cu129 | **Triton:** 3.5.1 > FLA baseline: [flash-linear-attention v0.5.0](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.5.0) +These results use cuLA's CuTe DSL FlashKDA backend. They supersede the earlier +numbers for the legacy CUTLASS C++ fused backend. +## FlashKDA Prefill vs FLA -## KDA Fused Forward (Kimi Delta Attention) - -Fully-fused KDA forward prefill kernel (sm90). - -### Fixed-Length (H=64, D=128, bf16) - -| B | T | FLA Triton (ms) | cuLA Fused (ms) | Speedup | -|---|---|-----------------|-----------------|---------| -| 1 | 512 | 0.556 | 0.224 | **2.48x** | -| 1 | 1024 | 0.581 | 0.248 | **2.34x** | -| 1 | 4096 | 0.936 | 0.896 | **1.04x** | -| 1 | 8192 | 1.810 | 1.754 | **1.03x** | -| 1 | 16384 | 3.576 | 3.492 | **1.02x** | -| 2 | 512 | 0.567 | 0.226 | **2.51x** | -| 2 | 1024 | 0.585 | 0.315 | **1.86x** | -| 2 | 4096 | 1.815 | 1.170 | **1.55x** | -| 2 | 8192 | 3.576 | 2.283 | **1.57x** | -| 2 | 16384 | 7.115 | 4.408 | **1.61x** | - -### Variable-Length (H=64, D=128, bf16) - -| Config | FLA Triton (ms) | cuLA Fused (ms) | Speedup | -|--------|-----------------|-----------------|---------| -| uniform 10seqs T=4096 [409..415] avg=409 | 1.019 | 0.707 | **1.44x** | -| random 10seqs T=4096 [24..1201] avg=409 | 1.013 | 0.669 | **1.51x** | -| skewed 10seqs T=4096 [227..2053] avg=409 | 1.010 | 0.681 | **1.48x** | -| uniform 20seqs T=4096 [204..220] avg=204 | 1.098 | 0.932 | **1.18x** | -| random 20seqs T=4096 [5..787] avg=204 | 1.074 | 0.748 | **1.44x** | -| skewed 20seqs T=4096 [107..2063] avg=204 | 1.048 | 0.732 | **1.43x** | -| uniform 10seqs T=8192 [819..821] avg=819 | 1.851 | 1.174 | **1.58x** | -| random 10seqs T=8192 [48..2401] avg=819 | 1.890 | 1.217 | **1.55x** | -| skewed 10seqs T=8192 [455..4097] avg=819 | 1.905 | 1.225 | **1.55x** | -| uniform 20seqs T=8192 [409..421] avg=409 | 1.960 | 1.406 | **1.39x** | -| random 20seqs T=8192 [9..1574] avg=409 | 1.953 | 1.290 | **1.51x** | -| skewed 20seqs T=8192 [215..4107] avg=409 | 1.957 | 1.300 | **1.51x** | -| uniform 10seqs T=16384 [1638..1642] avg=1638 | 3.642 | 2.162 | **1.68x** | -| random 10seqs T=16384 [95..4802] avg=1638 | 3.609 | 2.279 | **1.58x** | -| skewed 10seqs T=16384 [910..8194] avg=1638 | 3.625 | 2.354 | **1.54x** | -| uniform 20seqs T=16384 [819..823] avg=819 | 3.644 | 2.320 | **1.57x** | -| random 20seqs T=16384 [19..3147] avg=819 | 3.681 | 2.293 | **1.61x** | -| skewed 20seqs T=16384 [431..8195] avg=819 | 3.634 | 2.371 | **1.53x** | - -Summary (28 configs): **avg=1.58x**, min=1.02x, max=2.51x. +CuTe DSL FlashKDA runs a two-kernel K1+K2 prefill pipeline. The benchmark uses +BF16 inputs, `H=64`, `D=128`, `safe_gate=True`, no initial state, 25 warmup +iterations, and 100 measured iterations aggregated with the IQR mean. + +### Fixed-Length + +| B | T | FLA Triton (ms) | cuLA FlashKDA (ms) | Speedup | +|---|---:|---:|---:|---:| +| 1 | 512 | 1.0013 | 0.1284 | **7.80x** | +| 1 | 1024 | 0.9747 | 0.1832 | **5.32x** | +| 1 | 4096 | 0.9994 | 0.6537 | **1.53x** | +| 1 | 8192 | 1.7774 | 1.2831 | **1.39x** | +| 1 | 16384 | 3.5038 | 2.5421 | **1.38x** | +| 2 | 512 | 0.9855 | 0.1386 | **7.11x** | +| 2 | 1024 | 1.0283 | 0.2447 | **4.20x** | +| 2 | 4096 | 1.7690 | 0.9010 | **1.96x** | +| 2 | 8192 | 3.4866 | 1.7810 | **1.96x** | +| 2 | 16384 | 6.9413 | 3.5363 | **1.96x** | + +### Variable-Length + +| Config | FLA Triton (ms) | cuLA FlashKDA (ms) | Speedup | +|---|---:|---:|---:| +| uniform 10 seqs, T=4096, 409–415 tokens | 1.0295 | 0.4748 | **2.17x** | +| random 10 seqs, T=4096, 24–1201 tokens | 1.0378 | 0.5544 | **1.87x** | +| skewed 10 seqs, T=4096, 227–2053 tokens | 1.0126 | 0.6700 | **1.51x** | +| uniform 20 seqs, T=4096, 204–220 tokens | 1.0510 | 0.4639 | **2.27x** | +| random 20 seqs, T=4096, 5–787 tokens | 1.0250 | 0.5513 | **1.86x** | +| skewed 20 seqs, T=4096, 107–2063 tokens | 1.0048 | 0.6767 | **1.48x** | +| uniform 10 seqs, T=8192, 819–821 tokens | 1.7923 | 0.9051 | **1.98x** | +| random 10 seqs, T=8192, 48–2401 tokens | 1.8277 | 1.0525 | **1.74x** | +| skewed 10 seqs, T=8192, 455–4097 tokens | 1.8457 | 1.2728 | **1.45x** | +| uniform 20 seqs, T=8192, 409–421 tokens | 1.8855 | 0.8628 | **2.19x** | +| random 20 seqs, T=8192, 9–1574 tokens | 1.8785 | 1.0189 | **1.84x** | +| skewed 20 seqs, T=8192, 215–4107 tokens | 1.8875 | 1.2913 | **1.46x** | +| uniform 10 seqs, T=16384, 1638–1642 tokens | 3.5091 | 1.7505 | **2.00x** | +| random 10 seqs, T=16384, 95–4802 tokens | 3.5175 | 2.0473 | **1.72x** | +| skewed 10 seqs, T=16384, 910–8194 tokens | 3.5350 | 2.4892 | **1.42x** | +| uniform 20 seqs, T=16384, 819–823 tokens | 3.5335 | 1.6558 | **2.13x** | +| random 20 seqs, T=16384, 19–3147 tokens | 3.5694 | 1.9728 | **1.81x** | +| skewed 20 seqs, T=16384, 431–8195 tokens | 3.5403 | 2.4949 | **1.42x** | + +Across all 28 fixed-length and variable-length configs, FlashKDA averages +**2.39x** over FLA (minimum **1.38x**, maximum **7.80x**). To reproduce: ```bash python benchmarks/bench_kda_sm90_prefill.py --mode both ``` + +## Intracard Context Parallelism + +The following comparison uses the same FlashKDA binary with +`use_intracard_cp="auto"` and `use_intracard_cp=False`. It covers the long +single-sequence and ragged packed-sequence shapes where intracard CP fills the +H200's SM array more effectively. Each row uses BF16, `D=128`, +`safe_gate=True`, 10 warmup iterations, and 100 measured iterations. + +| Config | H | CP engaged | CP off (ms) | CP auto (ms) | Speedup | +|---|---:|:---:|---:|---:|---:| +| T=1023 | 4 | no | 0.1322 | 0.1447 | 0.91x | +| T=1025 | 4 | no | 0.1364 | 0.1388 | 0.98x | +| T=4K | 4 | yes | 0.4374 | 0.2163 | **2.02x** | +| T=8K | 4 | yes | 0.8636 | 0.2820 | **3.06x** | +| T=16K | 4 | yes | 1.7099 | 0.4078 | **4.19x** | +| T=32K | 4 | yes | 3.3888 | 0.5973 | **5.67x** | +| T=64K | 4 | yes | 6.7518 | 0.9658 | **6.99x** | +| T=64K+1 | 4 | yes | 7.3033 | 1.0126 | **7.21x** | +| 2x16K | 4 | yes | 1.7631 | 0.5436 | **3.24x** | +| 32K+4K | 4 | yes | 3.4000 | 0.6682 | **5.09x** | +| 32K+1K | 4 | yes | 3.3923 | 0.6046 | **5.61x** | +| 32K+1023+1025 | 4 | yes | 3.6605 | 0.6821 | **5.37x** | +| 64K+1K | 4 | yes | 6.7713 | 1.0644 | **6.36x** | +| 64K+2x1K | 4 | yes | 6.7412 | 1.0803 | **6.24x** | +| 64K+5x1K | 4 | yes | 6.7648 | 1.1064 | **6.11x** | +| 64K+1+1023+1025 | 4 | yes | 7.3010 | 1.1200 | **6.52x** | +| T=1023 | 8 | no | 0.1365 | 0.1395 | 0.98x | +| T=1025 | 8 | no | 0.1376 | 0.1391 | 0.99x | +| T=4K | 8 | yes | 0.4557 | 0.2356 | **1.93x** | +| T=8K | 8 | yes | 0.8941 | 0.3177 | **2.81x** | +| T=16K | 8 | yes | 1.7600 | 0.5132 | **3.43x** | +| T=32K | 8 | yes | 3.4958 | 0.8774 | **3.98x** | +| T=64K | 8 | yes | 6.9656 | 1.6239 | **4.29x** | +| T=64K+1 | 8 | yes | 7.5212 | 1.7221 | **4.37x** | +| 2x16K | 8 | yes | 1.8764 | 0.8328 | **2.25x** | +| 32K+4K | 8 | yes | 3.5256 | 1.0458 | **3.37x** | +| 32K+1K | 8 | yes | 3.5080 | 0.9871 | **3.55x** | +| 32K+1023+1025 | 8 | yes | 3.7879 | 1.0731 | **3.53x** | +| 64K+1K | 8 | yes | 6.9812 | 1.8345 | **3.81x** | +| 64K+2x1K | 8 | yes | 6.9297 | 1.8368 | **3.77x** | +| 64K+5x1K | 8 | yes | 6.9738 | 1.9318 | **3.61x** | +| 64K+1+1023+1025 | 8 | yes | 7.4816 | 1.9172 | **3.90x** | + +Intracard CP engages for 28 of the 32 configs. On the engaged subset it +delivers a **4.11x geometric-mean** speedup, with a range of **1.93–7.21x**. +The non-CHUNK-aligned `T=64K+1` cases retain essentially the same benefit as +`T=64K`; the four 1023/1025-token rows stay on the serial path and expose only +the auto-planner overhead. + +To reproduce: + +```bash +python benchmarks/bench_kda_sm90_cp.py +``` diff --git a/README.md b/README.md index 781006e..b154782 100644 --- a/README.md +++ b/README.md @@ -156,22 +156,26 @@ See [BENCHMARK_GB200_CUDA_130.md](BENCHMARK_GB200_CUDA_130.md) tested with CUDA **Hopper (SM90)** -See [BENCHMARK_H200.md](BENCHMARK_H200.md) tested with CUDA 12.9 for detailed results. +See [BENCHMARK_H200.md](BENCHMARK_H200.md) for CuTe DSL FlashKDA results on an H200 141GB with CUDA 12.9. **Highlights:** - **KDA Modular Forward (Blackwell):** **avg 1.33x** speedup on fixed-length, **avg 1.35x** on variable-length (18 configs, uniform/skewed/random). - **Lightning Attention Prefill (Blackwell):** up to **2.08x** speedup (B=2). - **Lightning Attention Varlen (Blackwell):** **avg 1.47x** speedup across 126 configs (uniform/skewed/random). -- **KDA Fused Forward (Hopper):** **avg 1.58x** speedup across fixed-length and variable-length sequences. +- **FlashKDA Prefill (Hopper):** **avg 2.39x** speedup over FLA across 28 fixed-length and variable-length configs, up to **7.80x**. +- **FlashKDA Intracard CP (Hopper):** **4.11x geo-mean** speedup over serial FlashKDA on 28 CP-engaged long-sequence configs, up to **7.21x**. -To regenerate benchmarks: +To reproduce the benchmark suites directly: ```bash # Blackwell (SM10X) -python benchmarks/generate_benchmark_md.py +python benchmarks/bench_kda.py --mode both +python benchmarks/bench_lightning_attn_prefill.py --modes no_state varlen +python benchmarks/bench_la_decode_vs_fla.py --heads 64 --head-dim 128 # Hopper (SM90) -python benchmarks/generate_benchmark_hopper_md.py +python benchmarks/bench_kda_sm90_prefill.py --mode both +python benchmarks/bench_kda_sm90_cp.py ``` ## Tests diff --git a/benchmarks/generate_benchmark_hopper_md.py b/benchmarks/generate_benchmark_hopper_md.py deleted file mode 100644 index cb9e8e0..0000000 --- a/benchmarks/generate_benchmark_hopper_md.py +++ /dev/null @@ -1,212 +0,0 @@ -#!/usr/bin/env python3 -""" -generate_benchmark_hopper_md.py — Run Hopper (SM90) benchmarks and generate BENCHMARK_hopper.md - -Currently supported Hopper benchmarks: - - KDA prefill (cula.kda.flashkda — SM90 K1+K2 two-kernel) - -Reuses bench_kda_sm90_prefill.py (calls cula_kda_prefill directly). - -Usage: - python benchmarks/generate_benchmark_hopper_md.py - - # Generate with specific GPU id - CUDA_VISIBLE_DEVICES=3 python benchmarks/generate_benchmark_hopper_md.py - - # Use cached results (skip re-running benchmarks) - python benchmarks/generate_benchmark_hopper_md.py --cache tmp/bench_hopper_cache.json -""" - -import argparse -import json -import os -import sys -from datetime import datetime -from pathlib import Path - -import numpy as np - -# Ensure project root is on path -ROOT = Path(__file__).resolve().parent.parent -sys.path.insert(0, str(ROOT)) - -os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) # Enable fast ops in FLA for fair comparison - -from benchmarks.bench_kda_sm90_prefill import ( # noqa: E402 - _SM_TAG, -) -from benchmarks.bench_kda_sm90_prefill import ( # noqa: E402 - D as KDA_D, -) -from benchmarks.bench_kda_sm90_prefill import ( # noqa: E402 - H as KDA_H, -) -from benchmarks.bench_kda_sm90_prefill import ( # noqa: E402 - main as kda_prefill_main, -) -from benchmarks.utils import get_env_info # noqa: E402 - -BENCHMARK_MD_DEFAULT = "BENCHMARK_H200.md" - - -# ============================================================ -# Run benchmarks -# ============================================================ - - -def run_kda_prefill_benchmarks(has_init_state: bool = False, heads=None): - """Run bench_kda_sm90_prefill.main() with programmatic args and return (fixed, varlen) results.""" - print("\n>>> Running SM90 KDA prefill benchmarks (via bench_kda_sm90_prefill.main)...") - orig_argv = sys.argv - argv = ["bench_kda_sm90_prefill.py", "--mode", "both"] - if has_init_state: - argv.append("--init_state") - if heads is not None: - argv += ["--heads", str(heads)] - sys.argv = argv - try: - fixed_res, varlen_res = kda_prefill_main() - finally: - sys.argv = orig_argv - return fixed_res, varlen_res - - -# ============================================================ -# Format markdown -# ============================================================ - - -def format_benchmark_md(env, kda_fixed, kda_varlen, has_init_state: bool = False): - lines = [] - w = lines.append - - w("# Benchmark Results — Hopper (SM90)\n") - w(f"> Auto-generated by `benchmarks/generate_benchmark_hopper_md.py` on {datetime.now().strftime('%Y-%m-%d')}.\n") - w(f"> **GPU:** {env['gpu']} | **CUDA:** {env['cuda']} | **PyTorch:** {env['torch']}\n") - w( - "> FLA baseline: [flash-linear-attention v0.5.0](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.5.0)\n" - ) - w("") - - # ------------------------------------------------------------------- - # KDA prefill - # ------------------------------------------------------------------- - w("\n## KDA Prefill (Kimi Delta Attention)\n") - w(f"SM90 K1+K2 two-kernel prefill ({_SM_TAG}).\n") - - # Fixed-length - if kda_fixed: - w(f"### Fixed-Length (H={KDA_H}, D={KDA_D}, bf16)\n") - w("| B | T | FLA Triton (ms) | cuLA (ms) | Speedup |") - w("|---|---|-----------------|-----------|---------|") - for r in kda_fixed: - sp = f"**{r['speedup']:.2f}x**" - w(f"| {r['B']} | {r['T']} | {r['ms_fla']:.3f} | {r['ms_cula']:.3f} | {sp} |") - - # Varlen - if kda_varlen: - w(f"\n### Variable-Length (H={KDA_H}, D={KDA_D}, bf16)\n") - w("| Config | FLA Triton (ms) | cuLA (ms) | Speedup |") - w("|--------|-----------------|-----------|---------|") - for r in kda_varlen: - tag = r.get("tag", "unknown") - sp = f"**{r['speedup']:.2f}x**" - w(f"| {tag} | {r['ms_fla']:.3f} | {r['ms_cula']:.3f} | {sp} |") - - # Summary stats - all_results = (kda_fixed or []) + (kda_varlen or []) - all_sp = [r["speedup"] for r in all_results if r.get("speedup", 0) > 0] - if all_sp: - w( - f"\nSummary ({len(all_sp)} configs): " - f"**avg={np.mean(all_sp):.2f}x**, min={np.min(all_sp):.2f}x, max={np.max(all_sp):.2f}x.\n" - ) - - w("To reproduce:\n") - w("```bash") - init_state_flag = " --init_state" if has_init_state else "" - w(f"python benchmarks/bench_kda_sm90_prefill.py --mode both{init_state_flag}") - w("```\n") - - return "\n".join(lines) - - -# ============================================================ -# Main -# ============================================================ - - -def main(): - parser = argparse.ArgumentParser(description="Generate BENCHMARK_hopper.md") - parser.add_argument( - "--cache", type=str, default=None, help="Path to a JSON cache file. If exists, skip benchmarks and use cached results." - ) - parser.add_argument( - "--output", - type=str, - default=None, - help="Output markdown filename (relative to project root). Default: BENCHMARK_H200.md", - ) - parser.add_argument("--save-cache", type=str, default=None, help="Save benchmark results to JSON for future --cache use.") - parser.add_argument( - "--init_state", - action="store_true", - help="Use non-zero initial state (default: False)", - ) - parser.add_argument( - "--heads", - type=int, - default=None, - help="Number of heads (H == HV, MHA) for KDA benchmarks. Default: use bench_kda_sm90_prefill default.", - ) - args = parser.parse_args() - - env = get_env_info() - - if args.cache and os.path.exists(args.cache): - print(f"Loading cached results from {args.cache}") - with open(args.cache) as f: - data = json.load(f) - kda_fixed = data["kda_fixed"] - kda_varlen = data["kda_varlen"] - else: - kda_fixed, kda_varlen = run_kda_prefill_benchmarks(has_init_state=args.init_state, heads=args.heads) - - if args.save_cache: - cache_path = Path(args.save_cache) - cache_path.parent.mkdir(parents=True, exist_ok=True) - - def sanitize(results): - out = [] - for r in results: - clean = {} - for k, v in r.items(): - if isinstance(v, (list, tuple)): - clean[k] = list(v) - elif isinstance(v, float) and (np.isnan(v) or np.isinf(v)): - clean[k] = None - else: - clean[k] = v - out.append(clean) - return out - - with open(cache_path, "w") as f: - json.dump( - { - "kda_fixed": sanitize(kda_fixed), - "kda_varlen": sanitize(kda_varlen), - }, - f, - indent=2, - ) - print(f"Cached results to {cache_path}") - - md = format_benchmark_md(env, kda_fixed, kda_varlen, has_init_state=args.init_state) - - output_path = ROOT / (args.output if args.output else BENCHMARK_MD_DEFAULT) - output_path.write_text(md) - print(f"\nWrote {output_path}") - - -if __name__ == "__main__": - main() diff --git a/benchmarks/generate_benchmark_md.py b/benchmarks/generate_benchmark_md.py deleted file mode 100644 index b9d9ff1..0000000 --- a/benchmarks/generate_benchmark_md.py +++ /dev/null @@ -1,368 +0,0 @@ -#!/usr/bin/env python3 -# Copyright 2025-2026 Ant Group Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -generate_benchmark_md.py — Run existing bench_* scripts and generate BENCHMARK.md - -Reuses bench_kda.py and bench_lightning_attn_prefill.py to collect results, -then formats them as BENCHMARK.md. - -Usage: - python benchmarks/generate_benchmark_md.py - - # Generate with specific GPU id - CUDA_VISIBLE_DEVICES=3 python benchmarks/generate_benchmark_md.py - - # Use cached results (skip re-running benchmarks) - python benchmarks/generate_benchmark_md.py --cache tmp/bench_cache.json -""" - -import argparse -import json -import os -import sys -from datetime import datetime -from pathlib import Path -from types import SimpleNamespace - -# Ensure project root is on path -ROOT = Path(__file__).resolve().parent.parent -sys.path.insert(0, str(ROOT)) - -os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) # Enable fast ops in FLA for fair comparison - -import numpy as np # noqa: E402 - -from benchmarks.bench_kda import ( # noqa: E402 - D as KDA_D, -) -from benchmarks.bench_kda import ( # noqa: E402 - H as KDA_H, -) -from benchmarks.bench_kda import ( # noqa: E402 - main as kda_main, -) -from benchmarks.bench_la_decode_vs_fla import ( # noqa: E402 - main as la_decode_main, -) -from benchmarks.bench_lightning_attn_prefill import ( # noqa: E402 - D_DEFAULT as LA_D, -) -from benchmarks.bench_lightning_attn_prefill import ( # noqa: E402 - _valid, -) -from benchmarks.bench_lightning_attn_prefill import ( # noqa: E402 - run_benchmark_suite as la_run_suite, -) -from benchmarks.utils import get_env_info # noqa: E402 - -BENCHMARK_MD_DEFAULT = "BENCHMARK_GB200.md" - -LA_H = 64 # for display only - - -# ============================================================ -# Run benchmarks by calling existing bench scripts -# ============================================================ - - -def run_kda_benchmarks(heads=None, hv=None): - """Run bench_kda.main() with programmatic args and return (fixed, varlen) results.""" - print("\n>>> Running KDA benchmarks (via bench_kda.main)...") - # bench_kda.main() parses sys.argv — override it temporarily - orig_argv = sys.argv - argv = ["bench_kda.py", "--mode", "both"] - if heads is not None: - argv += ["--heads", str(heads)] - if hv is not None: - argv += ["--hv", str(hv)] - sys.argv = argv - try: - fixed_res, varlen_res = kda_main() - finally: - sys.argv = orig_argv - return fixed_res, varlen_res - - -def run_la_decode_benchmarks(): - """Run bench_la_decode_vs_fla benchmarks with programmatic args.""" - print("\n>>> Running LA Decode benchmarks (via bench_la_decode_vs_fla.main)...") - orig_argv = sys.argv - sys.argv = [ - "bench_la_decode_vs_fla.py", - "--batch-sizes", - "1", - "4", - "16", - "64", - "256", - "--heads", - "64", - "--head-dim", - "128", - ] - try: - results = la_decode_main() - finally: - sys.argv = orig_argv - return results - - -def run_lightning_attn_benchmarks(): - """Run the Lightning Attention prefill benchmark with programmatic args.""" - print("\n>>> Running Lightning Attention prefill benchmarks...") - args = SimpleNamespace( - modes=["no_state", "varlen"], - batch_sizes=[1, 2], - seq_lens=[1024, 4096, 8192, 16384], - num_heads=[64], - head_dim=128, - layer_idx=12, - num_layers=24, - warmup=5, - iterations=20, - plot=False, - report=False, - ) - all_results = la_run_suite(args) - standard = [r for r in all_results if r["mode"] == "no_state"] - varlen = [r for r in all_results if r["mode"] == "varlen"] - return standard, varlen - - -# ============================================================ -# Format markdown -# ============================================================ - - -def format_benchmark_md(env, kda_fixed, kda_varlen, la_standard, la_varlen, la_decode=None): - lines = [] - w = lines.append - - w("# Benchmark Results\n") - w(f"> Auto-generated by `benchmarks/generate_benchmark_md.py` on {datetime.now().strftime('%Y-%m-%d')}.\n") - w(f"> **GPU:** {env['gpu']} | **CUDA:** {env['cuda']} | **PyTorch:** {env['torch']}\n") - w( - "> FLA baseline: [flash-linear-attention v0.5.0](https://github.com/fla-org/flash-linear-attention/releases/tag/v0.5.0)\n" - ) - w("") - - # ------------------------------------------------------------------- - # KDA - # ------------------------------------------------------------------- - w("\n## KDA (Kimi Delta Attention)\n") - - # Fixed-length - w(f"### Fixed-Length (H={KDA_H}, D={KDA_D}, bf16)\n") - w("| B | T | FLA Triton (ms) | cuLA (ms) | Speedup |") - w("|---|---|-----------------|-----------|---------|") - for r in kda_fixed: - sp = f"**{r['speedup']:.2f}x**" - w(f"| {r['B']} | {r['T']} | {r['ms_fla']:.3f} | {r['ms_cula']:.3f} | {sp} |") - - kda_fixed_sp = [r["speedup"] for r in kda_fixed if _valid(r.get("speedup", float("nan")))] - if kda_fixed_sp: - w( - f"\nSummary ({len(kda_fixed_sp)} configs): " - f"**avg={np.mean(kda_fixed_sp):.2f}x**, min={np.min(kda_fixed_sp):.2f}x, max={np.max(kda_fixed_sp):.2f}x.\n" - ) - - # Varlen - w(f"\n### Variable-Length (H={KDA_H}, D={KDA_D}, bf16)\n") - w("| Config | FLA Triton (ms) | cuLA (ms) | Speedup |") - w("|--------|-----------------|-----------|---------|") - for r in kda_varlen: - tag = r.get("tag", "unknown") - sp = f"**{r['speedup']:.2f}x**" - w(f"| {tag} | {r['ms_fla']:.3f} | {r['ms_cula']:.3f} | {sp} |") - - kda_varlen_sp = [r["speedup"] for r in kda_varlen if _valid(r.get("speedup", float("nan")))] - if kda_varlen_sp: - w( - f"\nSummary ({len(kda_varlen_sp)} configs): " - f"**avg={np.mean(kda_varlen_sp):.2f}x**, min={np.min(kda_varlen_sp):.2f}x, max={np.max(kda_varlen_sp):.2f}x.\n" - ) - - w("\nTo reproduce:\n") - w("```bash") - w("python benchmarks/bench_kda.py --mode both") - w("```\n") - - # ------------------------------------------------------------------- - # Lightning Attention - # ------------------------------------------------------------------- - w("## Lightning Attention\n") - - # Standard prefill - w(f"### Prefill (H={LA_H}, D={LA_D}, bf16)\n") - w("| B | T | FLA Triton (ms) | cuLA (ms) | Speedup |") - w("|---|---|-----------------|-----------|---------|") - for r in la_standard: - if not _valid(r.get("speedup", float("nan"))): - continue - sp = f"**{r['speedup']:.2f}x**" - w(f"| {r['B']} | {r['T']} | {r['fla_ms']:.3f} | {r['cutedsl_ms']:.3f} | {sp} |") - - # Varlen summary (select representative rows to avoid verbosity) - w(f"\n### Variable-Length (H={LA_H}, D={LA_D}, bf16)\n") - w("Persistent CuTe DSL kernel vs FLA Triton varlen.\n") - - # Pick a representative subset: one per (N, T) — use "uniform" dist only for the table - # but show summary stats across all dists - representative = [r for r in la_varlen if r.get("dist") == "uniform"] - w("| N (seqs) | T | cuLA (ms) | FLA Triton (ms) | Speedup |") - w("|----------|---|-----------|-----------------|---------|") - for r in representative: - p_ms = r.get("persistent_ms", float("nan")) - fla_ms = r.get("fla_varlen_ms", float("nan")) - sp = r.get("p_vs_fla_vl_speedup", float("nan")) - if not _valid(p_ms) or not _valid(fla_ms): - continue - sp_str = f"**{sp:.2f}x**" if _valid(sp) else "-" - w(f"| {r['B']} | {r['T']} | {p_ms:.3f} | {fla_ms:.3f} | {sp_str} |") - - # Summary stats across all varlen configs - all_sp = [r["p_vs_fla_vl_speedup"] for r in la_varlen if _valid(r.get("p_vs_fla_vl_speedup", float("nan")))] - if all_sp: - w( - f"\nSummary ({len(all_sp)} configs across uniform/skewed/random): " - f"**avg={np.mean(all_sp):.2f}x**, min={np.min(all_sp):.2f}x, max={np.max(all_sp):.2f}x.\n" - ) - - w("To reproduce:\n") - w("```bash") - w("python benchmarks/bench_lightning_attn_prefill.py --modes no_state varlen") - w("```\n") - - # ------------------------------------------------------------------- - # Lightning Attention Decode - # ------------------------------------------------------------------- - if la_decode: - w("### Decode (H=64, D=128, bf16, T=1)\n") - w("Single-token decode: la_decode (CuTe DSL) vs fla fused_recurrent (Triton).\n") - - # Kernel-only table - w("#### Kernel-Only\n") - w("| B | FLA Triton (ms) | cuLA (ms) | Speedup |") - w("|---|-----------------|-----------|---------|") - for r in la_decode: - sp = f"**{r['kernel_speedup']:.2f}x**" - w(f"| {r['B']} | {r['kernel_fla_ms']:.4f} | {r['kernel_cute_ms']:.4f} | {sp} |") - - # Wrapper table - w("\n#### Wrapper (Full Call Path)\n") - w("| B | FLA Triton (ms) | cuLA (ms) | Speedup |") - w("|---|-----------------|-----------|---------|") - for r in la_decode: - sp = f"**{r['wrap_speedup']:.2f}x**" - w(f"| {r['B']} | {r['wrap_fla_ms']:.4f} | {r['wrap_cute_ms']:.4f} | {sp} |") - - w("\nTo reproduce:\n") - w("```bash") - w("python benchmarks/bench_la_decode_vs_fla.py --heads 64 --head-dim 128") - w("```\n") - - return "\n".join(lines) - - -# ============================================================ -# Main -# ============================================================ - - -def main(): - parser = argparse.ArgumentParser(description="Generate BENCHMARK.md") - parser.add_argument( - "--cache", type=str, default=None, help="Path to a JSON cache file. If exists, skip benchmarks and use cached results." - ) - parser.add_argument( - "--output", - type=str, - default=None, - help="Output markdown filename (relative to project root). Default: BENCHMARK_GB200.md", - ) - parser.add_argument("--save-cache", type=str, default=None, help="Save benchmark results to JSON for future --cache use.") - parser.add_argument( - "--heads", - type=int, - default=None, - help="Number of Q/K heads (H) for KDA benchmarks. Default: use bench_kda default.", - ) - parser.add_argument( - "--hv", - type=int, - default=None, - help="Number of V heads (HV) for KDA benchmarks. For GVA, set HV > H with HV %% H == 0.", - ) - args = parser.parse_args() - - env = get_env_info() - - if args.cache and os.path.exists(args.cache): - print(f"Loading cached results from {args.cache}") - with open(args.cache) as f: - data = json.load(f) - kda_fixed = data["kda_fixed"] - kda_varlen = data["kda_varlen"] - la_standard = data["la_standard"] - la_varlen = data["la_varlen"] - la_decode = data.get("la_decode", []) - else: - kda_fixed, kda_varlen = run_kda_benchmarks(heads=args.heads, hv=args.hv) - la_standard, la_varlen = run_lightning_attn_benchmarks() - la_decode = run_la_decode_benchmarks() - - if args.save_cache: - cache_path = Path(args.save_cache) - cache_path.parent.mkdir(parents=True, exist_ok=True) - - # Convert results to JSON-serializable form - def sanitize(results): - out = [] - for r in results: - clean = {} - for k, v in r.items(): - if isinstance(v, (list, tuple)): - clean[k] = list(v) - elif isinstance(v, float) and (np.isnan(v) or np.isinf(v)): - clean[k] = None - else: - clean[k] = v - out.append(clean) - return out - - with open(cache_path, "w") as f: - json.dump( - { - "kda_fixed": sanitize(kda_fixed), - "kda_varlen": sanitize(kda_varlen), - "la_standard": sanitize(la_standard), - "la_varlen": sanitize(la_varlen), - "la_decode": sanitize(la_decode), - }, - f, - indent=2, - ) - print(f"Cached results to {cache_path}") - - md = format_benchmark_md(env, kda_fixed, kda_varlen, la_standard, la_varlen, la_decode) - - output_path = ROOT / (args.output if args.output else BENCHMARK_MD_DEFAULT) - output_path.write_text(md) - print(f"\nWrote {output_path}") - - -if __name__ == "__main__": - main()