@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / cf567cbc71

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Attention] Improve attention benchmarks: configs and profiling (#39336)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed Jun 12, 2026 at 16:24 UTC cf567cbc71a467d8479411062917e9190ee11376

15 files changed +573 -167

benchmarks/attention_benchmarks/README.md

+9 -13

@@ -108,7 +108,6 @@ python benchmark.py \

108 --backends flash triton flashinfer \

109 --batch-specs "q2k" "8q1s1k" "2q2k_32q1s1k" \

110 --num-layers 10 \

111 - --repeats 5 \

111 --output-csv results.csv

112 ```

113

@@ -164,14 +163,17 @@ python benchmark.py \

163 # Model configuration

164 --num-layers N # Number of layers

165 --head-dim N # Head dimension

166 +--v-head-dim N # Value head dimension (defaults to --head-dim)

167 --num-q-heads N # Query heads

168 --num-kv-heads N # KV heads

169 --block-size N # Block size

170 +--kv-lora-rank N # MLA KV LoRA rank

171 +--qk-nope-head-dim N # MLA non-RoPE QK head dim

172 +--qk-rope-head-dim N # MLA RoPE QK head dim

173

174 # Benchmark settings

175 --device DEVICE # Device (default: cuda:0)

173 ---repeats N # Repetitions

174 ---warmup-iters N # Warmup iterations

176 +--warmup-ms N # Warmup window in ms for triton do_bench

177 --profile-memory # Profile memory usage

178

179 # Parameter sweeps

@@ -211,8 +213,6 @@ config = BenchmarkConfig(

213 num_kv_heads=1,

214 block_size=128,

215 device="cuda:0",

214 - repeats=5,

215 - warmup_iters=3,

216 )

217

218 # CUTLASS MLA with specific num_kv_splits

@@ -253,14 +253,10 @@ formatter.save_json(results, "output.json")

253

254 ## Tips

255

256 -**1. Warmup matters** - Use `--warmup-iters 10` for stable results

256 +**1. Save results** - Always use `--output-csv` or `--output-json`

257

258 -**2. Multiple repeats** - Use `--repeats 20` for low variance

258 +**2. Test incrementally** - Start with `--num-layers 1`

259

260 -**3. Save results** - Always use `--output-csv` or `--output-json`

260 +**3. Extended grammar** - Leverage spec decode, chunked prefill patterns

261

262 -**4. Test incrementally** - Start with `--num-layers 1 --repeats 1`

263 -

264 -**5. Extended grammar** - Leverage spec decode, chunked prefill patterns

265 -

266 -**6. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values

262 +**4. Parameter sweeps** - Use `--sweep-param` and `--sweep-values` to find optimal values

benchmarks/attention_benchmarks/benchmark.py

+179 -39

@@ -26,6 +26,9 @@ Examples:

26 """

28 import argparse

29 +import os

30 +import shutil

31 +import subprocess

32 import sys

33 from dataclasses import replace

34 from pathlib import Path

@@ -83,13 +86,15 @@ def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:

86 else:

87 return run_standard_attention_benchmark(config)

88 except Exception as e:

89 + error_msg = str(e) or repr(e)

90 return BenchmarkResult(

91 config=config,

92 mean_time=float("inf"),

93 + median_time=float("inf"),

94 std_time=0,

95 min_time=float("inf"),

96 max_time=float("inf"),

92 - error=str(e),

97 + error=error_msg,

98 )

100

@@ -115,9 +120,12 @@ def run_model_parameter_sweep(

120 """

121 all_results = []

122

118 - console.print(

119 - f"[yellow]Model sweep mode: testing {sweep.param_name} = {sweep.values}[/]"

123 + sweep_desc = (

124 + f"{sweep.param_name} = {sweep.values}"

125 + if sweep.param_name

126 + else f"{len(sweep.values)} configurations"

127 )

128 + console.print(f"[yellow]Model sweep mode: testing {sweep_desc}[/]")

129

130 total = len(backends) * len(batch_specs) * len(sweep.values)

131

@@ -125,9 +133,9 @@ def run_model_parameter_sweep(

133 for backend in backends:

134 for spec in batch_specs:

135 for value in sweep.values:

128 - # Create config with modified model parameter

136 + # Create config with modified model parameter(s)

137 config_args = base_config_args.copy()

130 - config_args[sweep.param_name] = value

138 + sweep.apply(config_args, value)

139

140 # Create config with original backend for running

141 clean_config = BenchmarkConfig(

@@ -144,13 +152,21 @@ def run_model_parameter_sweep(

152 all_results.append(result)

153

154 if not result.success:

155 + err_label = (

156 + f"{sweep.param_name}={value}"

157 + if sweep.param_name

158 + else f"{value}"

159 + )

160 console.print(

148 - f"[red]Error {backend} {spec} {sweep.param_name}="

149 - f"{value}: {result.error}[/]"

161 + f"[red]Error {backend} {spec} {err_label}"

162 + f": {result.error}[/]"

163 )

164

165 pbar.update(1)

166

167 + if base_config_args.get("ncu_profile"):

168 + return all_results

169 +

170 # Display sweep results - create separate table for each parameter value

171 console.print("\n[bold green]Model Parameter Sweep Results:[/]")

172 formatter = ResultsFormatter(console)

@@ -184,7 +200,10 @@ def run_model_parameter_sweep(

200 )

201

202 for param_value in sorted_param_values:

187 - console.print(f"\n[bold cyan]{sweep.param_name} = {param_value}[/]")

203 + label = (

204 + f"{sweep.param_name} = {param_value}" if sweep.param_name else param_value

205 + )

206 + console.print(f"\n[bold cyan]{label}[/]")

207 param_results = by_param_value[param_value]

208

209 # Create modified results with original backend names

@@ -200,8 +219,9 @@ def run_model_parameter_sweep(

219 formatter.print_table(modified_results, backends, compare_to_fastest=True)

220

221 # Show optimal backend for each (param_value, batch_spec) combination

222 + sweep_name = sweep.param_name or "config"

223 console.print(

204 - f"\n[bold cyan]Optimal backend for each ({sweep.param_name}, batch_spec):[/]"

224 + f"\n[bold cyan]Optimal backend for each ({sweep_name}, batch_spec):[/]"

225 )

226

227 # Group by (param_value, batch_spec)

@@ -236,7 +256,10 @@ def run_model_parameter_sweep(

256 for param_value, spec in sorted_keys:

257 # Print header when param value changes

258 if param_value != current_param_value:

239 - console.print(f"\n [bold]{sweep.param_name}={param_value}:[/]")

259 + header = (

260 + f"{sweep.param_name}={param_value}" if sweep.param_name else param_value

261 + )

262 + console.print(f"\n [bold]{header}:[/]")

263 current_param_value = param_value

264

265 results = by_param_and_spec[(param_value, spec)]

@@ -322,6 +345,9 @@ def run_parameter_sweep(

345

346 pbar.update(1)

347

348 + if base_config_args.get("ncu_profile"):

349 + return all_results

350 +

351 # Display sweep results

352 console.print("\n[bold green]Sweep Results:[/]")

353 backend_labels = [sweep.get_label(b, v) for b in backends for v in sweep_values]

@@ -474,11 +500,35 @@ def main():

500 parser.add_argument("--num-q-heads", type=int, default=32, help="Query heads")

501 parser.add_argument("--num-kv-heads", type=int, default=8, help="KV heads")

502 parser.add_argument("--block-size", type=int, default=16, help="Block size")

503 + parser.add_argument(

504 + "--v-head-dim",

505 + type=int,

506 + default=None,

507 + help="Value head dimension (defaults to --head-dim if unset)",

508 + )

509 +

510 + # MLA-specific model dimensions

511 + parser.add_argument(

512 + "--kv-lora-rank", type=int, default=None, help="MLA KV LoRA rank"

513 + )

514 + parser.add_argument(

515 + "--qk-nope-head-dim", type=int, default=None, help="MLA non-RoPE QK head dim"

516 + )

517 + parser.add_argument(

518 + "--qk-rope-head-dim", type=int, default=None, help="MLA RoPE QK head dim"

519 + )

520

521 # Benchmark settings

522 parser.add_argument("--device", default="cuda:0", help="Device")

480 - parser.add_argument("--repeats", type=int, default=1, help="Repetitions")

481 - parser.add_argument("--warmup-iters", type=int, default=3, help="Warmup iterations")

523 + parser.add_argument(

524 + "--warmup-ms",

525 + type=int,

526 + default=None,

527 + help=(

528 + "Warmup window in ms for triton's do_bench (default: triton's own). "

529 + "Has no effect with CUDA graphs; pass --no-cuda-graphs to use it."

530 + ),

531 + )

532 parser.add_argument("--profile-memory", action="store_true", help="Profile memory")

533 parser.add_argument(

534 "--kv-cache-dtype",

@@ -491,10 +541,33 @@ def main():

541 action=argparse.BooleanOptionalAction,

542 default=True,

543 help=(

494 - "Launch kernels with CUDA graphs to eliminate CPU overhead"

495 - "in measurements (default: True)"

544 + "Use triton do_bench_cudagraph (True) or do_bench (False) "

545 + "for timing. CUDA graphs eliminate CPU launch overhead "

546 + "(default: True)"

547 + ),

548 + )

549 + parser.add_argument(

550 + "--num-splits",

551 + type=int,

552 + default=None,

553 + help="FlashAttention split-K factor (0=auto heuristic, 1=disabled, >1=force N)",

554 + )

555 + parser.add_argument(

556 + "--ncu-profile",

557 + action="store_true",

558 + default=False,

559 + help=(

560 + "Enable Nsight Compute profiling mode. Automatically wraps the "

561 + "script with ncu, capturing a profile with source correlation. "

562 + "Use --ncu-output to set the output file name."

563 ),

564 )

565 + parser.add_argument(

566 + "--ncu-output",

567 + type=str,

568 + default="profile",

569 + help="Output file name for ncu profile (default: 'profile').",

570 + )

571

572 # Parameter sweep (use YAML config for advanced sweeps)

573 parser.add_argument(

@@ -576,23 +649,28 @@ def main():

649 model = yaml_config["model"]

650 args.num_layers = model.get("num_layers", args.num_layers)

651 args.head_dim = model.get("head_dim", args.head_dim)

652 + args.v_head_dim = model.get("v_head_dim", args.v_head_dim)

653 args.num_q_heads = model.get("num_q_heads", args.num_q_heads)

654 args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads)

655 args.block_size = model.get("block_size", args.block_size)

656 + # MLA-specific dimensions

657 + args.kv_lora_rank = model.get("kv_lora_rank", args.kv_lora_rank)

658 + args.qk_nope_head_dim = model.get("qk_nope_head_dim", args.qk_nope_head_dim)

659 + args.qk_rope_head_dim = model.get("qk_rope_head_dim", args.qk_rope_head_dim)

660

661 # Benchmark settings (top-level keys)

662 if "device" in yaml_config:

663 args.device = yaml_config["device"]

586 - if "repeats" in yaml_config:

587 - args.repeats = yaml_config["repeats"]

588 - if "warmup_iters" in yaml_config:

589 - args.warmup_iters = yaml_config["warmup_iters"]

664 + if "warmup_ms" in yaml_config:

665 + args.warmup_ms = yaml_config["warmup_ms"]

666 if "profile_memory" in yaml_config:

667 args.profile_memory = yaml_config["profile_memory"]

668 if "kv_cache_dtype" in yaml_config:

669 args.kv_cache_dtype = yaml_config["kv_cache_dtype"]

670 if "cuda_graphs" in yaml_config:

671 args.cuda_graphs = yaml_config["cuda_graphs"]

672 + if "ncu_profile" in yaml_config:

673 + args.ncu_profile = yaml_config["ncu_profile"]

674

675 # Parameter sweep configuration

676 if "parameter_sweep" in yaml_config:

@@ -612,7 +690,7 @@ def main():

690 if "model_parameter_sweep" in yaml_config:

691 sweep_config = yaml_config["model_parameter_sweep"]

692 args.model_parameter_sweep = ModelParameterSweep(

615 - param_name=sweep_config["param_name"],

693 + param_name=sweep_config.get("param_name"),

694 values=sweep_config["values"],

695 label_format=sweep_config.get(

696 "label_format", "{backend}_{param_name}_{value}"

@@ -631,6 +709,32 @@ def main():

709

710 console.print()

711

712 + # Re-exec under ncu if --ncu-profile and not already inside ncu. This runs

713 + # after YAML processing so ncu_profile set via config file is honored.

714 + if args.ncu_profile and "_NCU_INNER" not in os.environ:

715 + ncu = shutil.which("ncu")

716 + if ncu is None:

717 + print("Error: 'ncu' not found in PATH", file=sys.stderr)

718 + sys.exit(1)

719 + cmd = [

720 + ncu,

721 + "--profile-from-start",

722 + "off",

723 + "--set",

724 + "full",

725 + "--import-source",

726 + "yes",

727 + "-o",

728 + args.ncu_output,

729 + sys.executable,

730 + *sys.argv,

731 + ]

732 + env = os.environ.copy()

733 + env["CUTE_DSL_LINEINFO"] = "1"

734 + env["_NCU_INNER"] = "1"

735 + print(f"Launching: {' '.join(cmd)}")

736 + sys.exit(subprocess.call(cmd, env=env))

737 +

738 # Handle CLI-based parameter sweep (if not from YAML)

739 if (

740 (not hasattr(args, "parameter_sweep") or args.parameter_sweep is None)

@@ -655,6 +759,18 @@ def main():

759 console.print(f"Batch specs: {', '.join(args.batch_specs)}")

760 console.print(f"KV cache dtype: {args.kv_cache_dtype}")

761 console.print(f"CUDA graphs: {args.cuda_graphs}")

762 + if args.warmup_ms is not None and args.cuda_graphs:

763 + console.print(

764 + "[yellow]Warning: --warmup-ms is ignored with CUDA graphs "

765 + "(do_bench_cudagraph warms up internally). Pass --no-cuda-graphs "

766 + "to use it.[/]"

767 + )

768 + if args.num_splits == 0 and args.cuda_graphs:

769 + console.print(

770 + "[yellow]Warning: --num-splits 0 (FA3 heuristic) is not CUDA-graph "

771 + "compatible and may fail or fall back. Pass --no-cuda-graphs or use "

772 + "--num-splits >=1.[/]"

773 + )

774 console.print()

775

776 init_workspace_manager(args.device)

@@ -662,6 +778,15 @@ def main():

778 # Run benchmarks

779 all_results = []

780

781 + # Under ncu profiling the kernels run only to be captured by the profiler;

782 + # timings are placeholder zeros, so the result tables and saved metrics are

783 + # skipped. The Nsight Compute report (--ncu-output) holds the real data.

784 + if args.ncu_profile:

785 + console.print(

786 + "[dim]ncu profiling enabled: result tables and saved metrics are "

787 + "skipped (timings are placeholder zeros).[/]"

788 + )

789 +

790 # Handle special mode: decode_vs_prefill comparison

791 if hasattr(args, "mode") and args.mode == "decode_vs_prefill":

792 console.print("[yellow]Mode: Decode vs Prefill pipeline comparison[/]")

@@ -708,11 +833,11 @@ def main():

833 num_kv_heads=args.num_kv_heads,

834 block_size=args.block_size,

835 device=args.device,

711 - repeats=args.repeats,

712 - warmup_iters=args.warmup_iters,

836 profile_memory=args.profile_memory,

837 kv_cache_dtype=args.kv_cache_dtype,

838 use_cuda_graphs=args.cuda_graphs,

839 + ncu_profile=args.ncu_profile,

840 + warmup_ms=args.warmup_ms,

841 )

842

843 # Add decode pipeline config

@@ -749,6 +874,7 @@ def main():

874 result = BenchmarkResult(

875 config=config,

876 mean_time=timing["mean"],

877 + median_time=timing.get("median", timing["mean"]),

878 std_time=timing["std"],

879 min_time=timing["min"],

880 max_time=timing["max"],

@@ -770,6 +896,7 @@ def main():

896 result = BenchmarkResult(

897 config=config,

898 mean_time=float("inf"),

899 + median_time=float("inf"),

900 std_time=0,

901 min_time=float("inf"),

902 max_time=float("inf"),

@@ -779,6 +906,9 @@ def main():

906

907 pbar.update(1)

908

909 + if args.ncu_profile:

910 + return

911 +

912 # Display decode vs prefill results

913 console.print("\n[bold green]Decode vs Prefill Results:[/]")

914

@@ -858,15 +988,20 @@ def main():

988 base_config_args = {

989 "num_layers": args.num_layers,

990 "head_dim": args.head_dim,

991 + "v_head_dim": args.v_head_dim,

992 "num_q_heads": args.num_q_heads,

993 "num_kv_heads": args.num_kv_heads,

994 "block_size": args.block_size,

995 "device": args.device,

865 - "repeats": args.repeats,

866 - "warmup_iters": args.warmup_iters,

996 "profile_memory": args.profile_memory,

997 "kv_cache_dtype": args.kv_cache_dtype,

998 "use_cuda_graphs": args.cuda_graphs,

999 + "ncu_profile": args.ncu_profile,

1000 + "warmup_ms": args.warmup_ms,

1001 + "num_splits": args.num_splits,

1002 + "kv_lora_rank": args.kv_lora_rank,

1003 + "qk_nope_head_dim": args.qk_nope_head_dim,

1004 + "qk_rope_head_dim": args.qk_rope_head_dim,

1005 }

1006 all_results = run_model_parameter_sweep(

1007 backends,

@@ -882,15 +1017,17 @@ def main():

1017 base_config_args = {

1018 "num_layers": args.num_layers,

1019 "head_dim": args.head_dim,

1020 + "v_head_dim": args.v_head_dim,

1021 "num_q_heads": args.num_q_heads,

1022 "num_kv_heads": args.num_kv_heads,

1023 "block_size": args.block_size,

1024 "device": args.device,

889 - "repeats": args.repeats,

890 - "warmup_iters": args.warmup_iters,

1025 "profile_memory": args.profile_memory,

1026 "kv_cache_dtype": args.kv_cache_dtype,

1027 "use_cuda_graphs": args.cuda_graphs,

1028 + "ncu_profile": args.ncu_profile,

1029 + "warmup_ms": args.warmup_ms,

1030 + "num_splits": args.num_splits,

1031 }

1032 all_results = run_parameter_sweep(

1033 backends, args.batch_specs, base_config_args, args.parameter_sweep, console

@@ -914,15 +1051,17 @@ def main():

1051 batch_spec=spec,

1052 num_layers=args.num_layers,

1053 head_dim=args.head_dim,

1054 + v_head_dim=getattr(args, "v_head_dim", None),

1055 num_q_heads=args.num_q_heads,

1056 num_kv_heads=args.num_kv_heads,

1057 block_size=args.block_size,

1058 device=args.device,

921 - repeats=args.repeats,

922 - warmup_iters=args.warmup_iters,

1059 profile_memory=args.profile_memory,

1060 kv_cache_dtype=args.kv_cache_dtype,

1061 use_cuda_graphs=args.cuda_graphs,

1062 + ncu_profile=args.ncu_profile,

1063 + warmup_ms=args.warmup_ms,

1064 + num_splits=args.num_splits,

1065 )

1066

1067 result = run_benchmark(config)

@@ -935,9 +1074,10 @@ def main():

1074

1075 pbar.update(1)

1076

938 - console.print("\n[bold green]Results:[/]")

939 - formatter = ResultsFormatter(console)

940 - formatter.print_table(decode_results, backends)

1077 + if not args.ncu_profile:

1078 + console.print("\n[bold green]Results:[/]")

1079 + formatter = ResultsFormatter(console)

1080 + formatter.print_table(decode_results, backends)

1081

1082 # Run prefill backend comparison

1083 if prefill_backends:

@@ -962,9 +1102,8 @@ def main():

1102 num_kv_heads=args.num_kv_heads,

1103 block_size=args.block_size,

1104 device=args.device,

965 - repeats=args.repeats,

966 - warmup_iters=args.warmup_iters,

1105 profile_memory=args.profile_memory,

1106 + warmup_ms=args.warmup_ms,

1107 prefill_backend=pb,

1108 )

1109

@@ -980,16 +1119,17 @@ def main():

1119

1120 pbar.update(1)

1121

983 - console.print("\n[bold green]Prefill Backend Results:[/]")

984 - formatter = ResultsFormatter(console)

985 - formatter.print_table(

986 - prefill_results, prefill_backends, compare_to_fastest=True

987 - )

1122 + if not args.ncu_profile:

1123 + console.print("\n[bold green]Prefill Backend Results:[/]")

1124 + formatter = ResultsFormatter(console)

1125 + formatter.print_table(

1126 + prefill_results, prefill_backends, compare_to_fastest=True

1127 + )

1128

1129 all_results = decode_results + prefill_results

1130

991 - # Save results

992 - if all_results:

1131 + # Save results (skip ncu profiling runs: timings are placeholder zeros)

1132 + if all_results and not args.ncu_profile:

1133 formatter = ResultsFormatter(console)

1134 if args.output_csv:

1135 formatter.save_csv(all_results, args.output_csv)

benchmarks/attention_benchmarks/common.py

+54 -6

@@ -15,6 +15,8 @@ from batch_spec import get_batch_type, parse_batch_spec

15 from rich.console import Console

16 from rich.table import Table

18 +from vllm.triton_utils import triton

19 +

21 def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:

22 """

@@ -34,6 +36,30 @@ def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:

36 return (0, 0, 0)

39 +def run_do_bench(

40 + benchmark_fn,

41 + use_cuda_graphs: bool,

42 + warmup_ms: int | None = None,

43 +) -> list[float]:

44 + kwargs: dict[str, Any] = {"return_mode": "all"}

45 + if use_cuda_graphs:

46 + result = triton.testing.do_bench_cudagraph(benchmark_fn, **kwargs)

47 + else:

48 + if warmup_ms is not None:

49 + kwargs["warmup"] = warmup_ms

50 + result = triton.testing.do_bench(benchmark_fn, **kwargs)

51 + return result

52 +

53 +

54 +def run_ncu_profile(benchmark_fn) -> None:

55 + benchmark_fn()

56 + torch.accelerator.synchronize()

57 + torch.cuda.cudart().cudaProfilerStart()

58 + benchmark_fn()

59 + torch.accelerator.synchronize()

60 + torch.cuda.cudart().cudaProfilerStop()

61 +

62 +

63 # Mock classes for vLLM attention infrastructure

@@ -182,18 +208,37 @@ class ParameterSweep:

208

209 @dataclass

210 class ModelParameterSweep:

185 - """Configuration for sweeping a model configuration parameter."""

211 + """Configuration for sweeping model configuration parameter(s).

212

187 - param_name: str # Name of the model config parameter to sweep (e.g., "num_q_heads")

188 - values: list[Any] # List of values to test

189 - label_format: str = "{backend}_{param_name}_{value}" # Result label template

213 + Supports two modes:

214 + - Single param: param_name="head_dim", values=[128, 256, 512]

215 + - Multi param: values=[{head_dim: 192, v_head_dim: 128}, {head_dim: 256}]

216 + When values are dicts, each dict's keys are applied as config overrides.

217 + """

218 +

219 + param_name: str | None = None

220 + values: list[Any] | None = None

221 + label_format: str = "{backend}_{param_name}_{value}"

222

223 def get_label(self, backend: str, value: Any) -> str:

224 """Generate a label for a specific parameter value."""

225 + if isinstance(value, dict):

226 + return self.label_format.format(

227 + backend=backend, param_name=self.param_name, value=value, **value

228 + )

229 return self.label_format.format(

230 backend=backend, param_name=self.param_name, value=value

231 )

232

233 + def apply(self, config_args: dict, value: Any) -> None:

234 + """Apply a sweep value to config args."""

235 + if isinstance(value, dict):

236 + config_args.update(value)

237 + elif self.param_name is not None:

238 + config_args[self.param_name] = value

239 + else:

240 + raise ValueError("param_name must be set if sweep values are not dicts")

241 +

242

243 @dataclass

244 class BenchmarkConfig:

@@ -208,10 +253,10 @@ class BenchmarkConfig:

253 block_size: int

254 device: str

255 dtype: torch.dtype = torch.float16

211 - repeats: int = 1

212 - warmup_iters: int = 3

256 profile_memory: bool = False

257 use_cuda_graphs: bool = False

258 + ncu_profile: bool = False

259 + warmup_ms: int | None = None

260

261 # "auto" or "fp8"

262 kv_cache_dtype: str = "auto"

@@ -226,6 +271,7 @@ class BenchmarkConfig:

271 # Backend-specific tuning

272 num_kv_splits: int | None = None # CUTLASS MLA

273 reorder_batch_threshold: int | None = None # FlashAttn MLA, FlashMLA

274 + num_splits: int | None = None # FlashAttention split-K (0=auto, 1=disabled)

275

276

277 @dataclass

@@ -234,6 +280,7 @@ class BenchmarkResult:

280

281 config: BenchmarkConfig

282 mean_time: float # seconds

283 + median_time: float # seconds

284 std_time: float # seconds

285 min_time: float # seconds

286 max_time: float # seconds

@@ -252,6 +299,7 @@ class BenchmarkResult:

299 return {

300 "config": asdict(self.config),

301 "mean_time": self.mean_time,

302 + "median_time": self.median_time,

303 "std_time": self.std_time,

304 "min_time": self.min_time,

305 "max_time": self.max_time,

benchmarks/attention_benchmarks/configs/mla_decode.yaml

-2

@@ -56,8 +56,6 @@ backends:

56 - TOKENSPEED_MLA # Blackwell + R1 dims + FP8 KV (use --kv-cache-dtype fp8)

58 device: "cuda:0"

59 -repeats: 100

60 -warmup_iters: 10

59 profile_memory: true

61 # Backend-specific tuning

benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml

-2

@@ -51,8 +51,6 @@ backends:

51 - FLASHMLA # Hopper only

53 device: "cuda:0"

54 -repeats: 5

55 -warmup_iters: 3

54 profile_memory: true

56 # Analyze chunked prefill workspace size impact

benchmarks/attention_benchmarks/configs/mla_prefill.yaml

-2

@@ -124,5 +124,3 @@ prefill_backends:

124 - tokenspeed

125

126 device: "cuda:0"

127 -repeats: 20

128 -warmup_iters: 5

benchmarks/attention_benchmarks/configs/mla_sparse_decode.yaml

-2

@@ -53,6 +53,4 @@ backends:

53 - FLASHINFER_MLA_SPARSE

55 device: "cuda:0"

56 -repeats: 100

57 -warmup_iters: 10

56 profile_memory: true

benchmarks/attention_benchmarks/configs/mla_sparse_prefill.yaml

-2

@@ -57,6 +57,4 @@ backends:

57 - FLASHINFER_MLA_SPARSE

59 device: "cuda:0"

60 -repeats: 10

61 -warmup_iters: 3

60 profile_memory: true

benchmarks/attention_benchmarks/configs/reorder_threshold.yaml

-2

@@ -63,8 +63,6 @@ model:

64 # Benchmark settings

65 device: "cuda:0"

66 -repeats: 15 # More repeats for spec decode variance

67 -warmup_iters: 5

66 profile_memory: false

68 # Output

benchmarks/attention_benchmarks/configs/speculative_decode.yaml

-2

@@ -49,8 +49,6 @@ backends:

50 # Benchmark settings

51 device: "cuda:0"

52 -repeats: 10 # More repeats for statistical significance

53 -warmup_iters: 5

52 profile_memory: false

54 # Test these threshold values for optimization

benchmarks/attention_benchmarks/configs/standard_attention.yaml

-2

@@ -43,6 +43,4 @@ backends:

43 - FLASHINFER

45 device: "cuda:0"

46 -repeats: 5

47 -warmup_iters: 3

46 profile_memory: false

benchmarks/attention_benchmarks/configs/standard_decode.yaml new

+142

@@ -0,0 +1,142 @@

1 +# Standard attention decode benchmark configuration

2 +# Sweeps num_q_heads and num_kv_heads to isolate effects of:

3 +# 1. GQA ratio (fixed num_q_heads=32, vary num_kv_heads)

4 +# 2. Absolute head count (fixed 4:1 ratio, vary scale)

5 +

6 +model:

7 + num_layers: 32

8 + num_q_heads: 32 # Base value, overridden by sweep

9 + num_kv_heads: 8 # Base value, overridden by sweep

10 + head_dim: 128

11 + block_size: 16

12 +

13 +# Head count sweep: each entry overrides num_q_heads, num_kv_heads, and

14 +# head_dim where it differs from the base (128). Head counts are per-GPU

15 +# (i.e. after TP sharding).

16 +#

17 +# Group A — vary GQA ratio (fixed q=32, head_dim=128):

18 +# 32:32 (MHA), 32:8 (GQA 4:1), 32:4 (GQA 8:1), 32:1 (MQA)

19 +#

20 +# Groups B-E — real model configs at various TP degrees:

21 +# Model head_dim Full TP2 TP4 TP8

22 +# Llama 3 8B 128 32:8 16:4 8:2 4:1

23 +# Llama 3 70B 128 64:8 32:4 16:2 8:1

24 +# GPT-OSS 120B 64 64:8 32:4 16:2 8:1

25 +# Llama 3 405B 128 128:8 64:4 32:2 16:1

26 +model_parameter_sweep:

27 + values:

28 + # --- head_dim=128 (Llama 3 family) ---

29 + - { num_q_heads: 32, num_kv_heads: 32, head_dim: 128 } # MHA 1:1

30 + - { num_q_heads: 32, num_kv_heads: 1, head_dim: 128 } # MQA 32:1

31 + - { num_q_heads: 4, num_kv_heads: 1, head_dim: 128 } # Llama 3 8B TP8

32 + - { num_q_heads: 8, num_kv_heads: 2, head_dim: 128 } # Llama 3 8B TP4

33 + - { num_q_heads: 16, num_kv_heads: 4, head_dim: 128 } # Llama 3 8B TP2

34 + - { num_q_heads: 32, num_kv_heads: 8, head_dim: 128 } # Llama 3 8B TP1 / GQA 4:1

35 + - { num_q_heads: 8, num_kv_heads: 1, head_dim: 128 } # Llama 3 70B TP8

36 + - { num_q_heads: 16, num_kv_heads: 2, head_dim: 128 } # Llama 3 70B TP4

37 + - { num_q_heads: 32, num_kv_heads: 4, head_dim: 128 } # Llama 3 70B TP2 / GQA 8:1

38 + - { num_q_heads: 64, num_kv_heads: 8, head_dim: 128 } # Llama 3 70B TP1

39 + - { num_q_heads: 16, num_kv_heads: 1, head_dim: 128 } # Llama 3 405B TP8

40 + - { num_q_heads: 32, num_kv_heads: 2, head_dim: 128 } # Llama 3 405B TP4

41 + - { num_q_heads: 64, num_kv_heads: 4, head_dim: 128 } # Llama 3 405B TP2

42 + - { num_q_heads: 128, num_kv_heads: 8, head_dim: 128 } # Llama 3 405B TP1

43 + # --- head_dim=64 (GPT-OSS 120B) ---

44 + - { num_q_heads: 8, num_kv_heads: 1, head_dim: 64 } # GPT-OSS 120B TP8

45 + - { num_q_heads: 16, num_kv_heads: 2, head_dim: 64 } # GPT-OSS 120B TP4

46 + - { num_q_heads: 32, num_kv_heads: 4, head_dim: 64 } # GPT-OSS 120B TP2

47 + - { num_q_heads: 64, num_kv_heads: 8, head_dim: 64 } # GPT-OSS 120B TP1

48 + label_format: "{backend}_q{num_q_heads}kv{num_kv_heads}d{head_dim}"

49 +

50 +batch_specs:

51 + # ---- batch_size x seq_len grid (decode: q_len=1) ----

52 + # Small grid for quick iteration. Uncomment for full sweep.

53 +

54 + # Batch size 1

55 + - "q1s1k"

56 + - "q1s512"

57 + - "q1s2k"

58 + - "q1s4k"

59 + - "q1s8k"

60 + - "q1s16k"

61 + - "q1s32k"

62 +

63 + # Batch size 2

64 + - "2q1s512"

65 + - "2q1s1k"

66 + - "2q1s2k"

67 + - "2q1s4k"

68 + - "2q1s8k"

69 + - "2q1s16k"

70 + - "2q1s32k"

71 +

72 + # Batch size 4

73 + - "4q1s512"

74 + - "4q1s1k"

75 + - "4q1s2k"

76 + - "4q1s4k"

77 + - "4q1s8k"

78 + - "4q1s16k"

79 + - "4q1s32k"

80 +

81 + # Batch size 8

82 + - "8q1s1k"

83 + - "8q1s512"

84 + - "8q1s2k"

85 + - "8q1s4k"

86 + - "8q1s8k"

87 + - "8q1s16k"

88 + - "8q1s32k"

89 +

90 + # Batch size 16

91 + - "16q1s512"

92 + - "16q1s1k"

93 + - "16q1s2k"

94 + - "16q1s4k"

95 + - "16q1s8k"

96 + - "16q1s16k"

97 + - "16q1s32k"

98 +

99 + # Batch size 32

100 + - "32q1s512"

101 + - "32q1s1k"

102 + - "32q1s2k"

103 + - "32q1s4k"

104 + - "32q1s8k"

105 + - "32q1s16k"

106 + - "32q1s32k"

107 +

108 + # Batch size 64

109 + - "64q1s1k"

110 + - "64q1s512"

111 + - "64q1s2k"

112 + - "64q1s4k"

113 + - "64q1s8k"

114 + - "64q1s16k"

115 + - "64q1s32k"

116 +

117 + # Batch size 128

118 + - "128q1s512"

119 + - "128q1s1k"

120 + - "128q1s2k"

121 + - "128q1s4k"

122 + - "128q1s8k"

123 + - "128q1s16k"

124 + - "128q1s32k"

125 +

126 + # Batch size 256

127 + - "256q1s1k"

128 + - "256q1s512"

129 + - "256q1s2k"

130 + - "256q1s4k"

131 + - "256q1s8k"

132 + - "256q1s16k"

133 + - "256q1s32k"

134 +

135 +# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER

136 +backends:

137 + - FLASH_ATTN

138 + - TRITON_ATTN

139 + - FLASHINFER

140 +

141 +device: "cuda:0"

142 +profile_memory: false

benchmarks/attention_benchmarks/configs/standard_prefill.yaml new

+108

@@ -0,0 +1,108 @@

1 +# Standard attention prefill benchmark configuration