mirror of
https://github.com/vllm-project/vllm.git
synced 2026-10-28 09:45:26 +08:00
JartXandtjtanaa
00b6c4c565
Signed-off-by: JartX <sagformas@epdcenter.es> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>
1457 lines
56 KiB
Python
1457 lines
56 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Universal vLLM Attention Benchmark.
|
|
|
|
Benchmark any attention backend with the extended grammar.
|
|
Supports standard attention (Flash/Triton/FlashInfer) and MLA backends.
|
|
|
|
Examples:
|
|
# Standard attention
|
|
python benchmark.py --backends flash flashinfer --batch-specs "q2k" "8q1s1k"
|
|
|
|
# MLA backends
|
|
python benchmark.py --backends cutlass_mla flashinfer_mla --batch-specs "64q1s1k"
|
|
|
|
# Parameter sweep (CLI)
|
|
python benchmark.py --backend cutlass_mla \
|
|
--batch-specs "64q1s1k" \
|
|
--sweep-param num_kv_splits \
|
|
--sweep-values 1 4 8 16
|
|
|
|
# Parameter sweep (YAML config - recommended)
|
|
python benchmark.py --config configs/cutlass_numsplits.yaml
|
|
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from dataclasses import replace
|
|
from pathlib import Path
|
|
from typing import get_args
|
|
|
|
import yaml
|
|
from rich.console import Console
|
|
from tqdm import tqdm
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
|
|
|
|
from batch_spec import parse_batch_spec
|
|
from common import (
|
|
BenchmarkConfig,
|
|
BenchmarkResult,
|
|
ModelParameterSweep,
|
|
ParameterSweep,
|
|
ResultsFormatter,
|
|
batch_spec_sort_key,
|
|
is_mla_backend,
|
|
)
|
|
|
|
from vllm.config.cache import CacheDType
|
|
from vllm.v1.worker.workspace import init_workspace_manager
|
|
|
|
|
|
def _str2bool(v) -> bool:
|
|
if isinstance(v, bool):
|
|
return v
|
|
if v.lower() in ("true", "1", "yes", "t"):
|
|
return True
|
|
if v.lower() in ("false", "0", "no", "f"):
|
|
return False
|
|
raise argparse.ArgumentTypeError(f"expected a boolean, got {v!r}")
|
|
|
|
|
|
def run_standard_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:
|
|
"""Run standard attention benchmark (Flash/Triton/FlashInfer)."""
|
|
from runner import run_attention_benchmark
|
|
|
|
return run_attention_benchmark(config)
|
|
|
|
|
|
def run_mla_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
|
|
"""Run MLA benchmark with appropriate backend."""
|
|
from mla_runner import run_mla_benchmark as run_mla
|
|
|
|
return run_mla(
|
|
config.backend,
|
|
config,
|
|
prefill_backend=config.prefill_backend,
|
|
sparse_mla_force_mqa=config.sparse_mla_force_mqa,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
|
|
"""Run a single benchmark with proper backend selection.
|
|
|
|
Args:
|
|
config: BenchmarkConfig with backend, batch_spec, and model params
|
|
**kwargs: Additional arguments passed to MLA benchmarks
|
|
|
|
Returns:
|
|
BenchmarkResult (may have error field set on failure)
|
|
|
|
"""
|
|
try:
|
|
if is_mla_backend(config.backend):
|
|
return run_mla_benchmark(config, **kwargs)
|
|
else:
|
|
return run_standard_attention_benchmark(config)
|
|
except Exception as e:
|
|
error_msg = str(e) or repr(e)
|
|
return BenchmarkResult(
|
|
config=config,
|
|
mean_time=float("inf"),
|
|
median_time=float("inf"),
|
|
std_time=0,
|
|
min_time=float("inf"),
|
|
max_time=float("inf"),
|
|
error=error_msg,
|
|
)
|
|
|
|
|
|
def run_model_parameter_sweep(
|
|
backends: list[str],
|
|
batch_specs: list[str],
|
|
base_config_args: dict,
|
|
sweep: ModelParameterSweep,
|
|
console: Console,
|
|
) -> list[BenchmarkResult]:
|
|
"""Run model parameter sweep for given backends and batch specs.
|
|
|
|
Args:
|
|
backends: List of backend names
|
|
batch_specs: List of batch specifications
|
|
base_config_args: Base configuration arguments (num_layers, head_dim, etc.)
|
|
sweep: ModelParameterSweep configuration
|
|
console: Rich console for output
|
|
|
|
Returns:
|
|
List of BenchmarkResult objects
|
|
|
|
"""
|
|
all_results = []
|
|
|
|
sweep_desc = (
|
|
f"{sweep.param_name} = {sweep.values}"
|
|
if sweep.param_name
|
|
else f"{len(sweep.values)} configurations"
|
|
)
|
|
console.print(f"[yellow]Model sweep mode: testing {sweep_desc}[/]")
|
|
|
|
total = len(backends) * len(batch_specs) * len(sweep.values)
|
|
|
|
with tqdm(total=total, desc="Benchmarking") as pbar:
|
|
for backend in backends:
|
|
for spec in batch_specs:
|
|
for value in sweep.values:
|
|
# Create config with modified model parameter(s)
|
|
config_args = base_config_args.copy()
|
|
sweep.apply(config_args, value)
|
|
|
|
# Create config with original backend for running
|
|
clean_config = BenchmarkConfig(
|
|
backend=backend, batch_spec=spec, **config_args
|
|
)
|
|
|
|
# Run benchmark
|
|
result = run_benchmark(clean_config)
|
|
|
|
# Replace backend with labeled version for display
|
|
backend_label = sweep.get_label(backend, value)
|
|
labeled_config = replace(result.config, backend=backend_label)
|
|
result = replace(result, config=labeled_config)
|
|
all_results.append(result)
|
|
|
|
if not result.success:
|
|
err_label = (
|
|
f"{sweep.param_name}={value}"
|
|
if sweep.param_name
|
|
else f"{value}"
|
|
)
|
|
console.print(
|
|
f"[red]Error {backend} {spec} {err_label}"
|
|
f": {result.error}[/]"
|
|
)
|
|
|
|
pbar.update(1)
|
|
|
|
if base_config_args.get("ncu_profile"):
|
|
return all_results
|
|
|
|
# Display sweep results - create separate table for each parameter value
|
|
console.print("\n[bold green]Model Parameter Sweep Results:[/]")
|
|
formatter = ResultsFormatter(console)
|
|
|
|
# Group results by parameter value and extract backend mapping
|
|
by_param_value = {}
|
|
backend_mapping = {} # Maps labeled backend -> original backend
|
|
|
|
for r in all_results:
|
|
# Extract original backend and param value from labeled backend
|
|
# The label format is: {backend}_{param_name}_{value}
|
|
# We need to reverse engineer this
|
|
labeled_backend = r.config.backend
|
|
|
|
# Try each backend to find which one this result belongs to
|
|
for backend in backends:
|
|
for value in sweep.values:
|
|
expected_label = sweep.get_label(backend, value)
|
|
if labeled_backend == expected_label:
|
|
backend_mapping[labeled_backend] = backend
|
|
param_value = str(value)
|
|
|
|
if param_value not in by_param_value:
|
|
by_param_value[param_value] = []
|
|
by_param_value[param_value].append(r)
|
|
break
|
|
|
|
# Create a table for each parameter value
|
|
sorted_param_values = sorted(
|
|
by_param_value.keys(), key=lambda x: int(x) if x.isdigit() else x
|
|
)
|
|
|
|
for param_value in sorted_param_values:
|
|
label = (
|
|
f"{sweep.param_name} = {param_value}" if sweep.param_name else param_value
|
|
)
|
|
console.print(f"\n[bold cyan]{label}[/]")
|
|
param_results = by_param_value[param_value]
|
|
|
|
# Create modified results with original backend names
|
|
modified_results = []
|
|
for r in param_results:
|
|
# Get the original backend name from our mapping
|
|
original_backend = backend_mapping[r.config.backend]
|
|
modified_config = replace(r.config, backend=original_backend)
|
|
modified_result = replace(r, config=modified_config)
|
|
modified_results.append(modified_result)
|
|
|
|
# Print table with original backend names
|
|
formatter.print_table(modified_results, backends, compare_to_fastest=True)
|
|
|
|
# Show optimal backend for each (param_value, batch_spec) combination
|
|
sweep_name = sweep.param_name or "config"
|
|
console.print(
|
|
f"\n[bold cyan]Optimal backend for each ({sweep_name}, batch_spec):[/]"
|
|
)
|
|
|
|
# Group by (param_value, batch_spec)
|
|
by_param_and_spec = {}
|
|
for r in all_results:
|
|
if r.success:
|
|
# Find which (backend, value) this result corresponds to
|
|
labeled_backend = r.config.backend
|
|
for backend in backends:
|
|
for value in sweep.values:
|
|
expected_label = sweep.get_label(backend, value)
|
|
if labeled_backend == expected_label:
|
|
param_value = str(value)
|
|
spec = r.config.batch_spec
|
|
key = (param_value, spec)
|
|
|
|
if key not in by_param_and_spec:
|
|
by_param_and_spec[key] = []
|
|
by_param_and_spec[key].append(r)
|
|
break
|
|
|
|
# Sort by param value then spec (batch_size, q_len, kv_len)
|
|
sorted_keys = sorted(
|
|
by_param_and_spec.keys(),
|
|
key=lambda x: (
|
|
int(x[0]) if x[0].isdigit() else x[0],
|
|
batch_spec_sort_key(x[1]),
|
|
),
|
|
)
|
|
|
|
current_param_value = None
|
|
for param_value, spec in sorted_keys:
|
|
# Print header when param value changes
|
|
if param_value != current_param_value:
|
|
header = (
|
|
f"{sweep.param_name}={param_value}" if sweep.param_name else param_value
|
|
)
|
|
console.print(f"\n [bold]{header}:[/]")
|
|
current_param_value = param_value
|
|
|
|
results = by_param_and_spec[(param_value, spec)]
|
|
best = min(results, key=lambda r: r.mean_time)
|
|
|
|
# Extract original backend name using the mapping
|
|
backend_name = backend_mapping[best.config.backend]
|
|
|
|
# Show all backends' times for comparison
|
|
times_str = " | ".join(
|
|
[
|
|
f"{backend_mapping[r.config.backend]}: {r.mean_time:.6f}s"
|
|
for r in sorted(results, key=lambda r: r.mean_time)
|
|
]
|
|
)
|
|
|
|
console.print(
|
|
f" {spec:12s} -> [bold green]{backend_name:15s}[/] ({times_str})"
|
|
)
|
|
|
|
return all_results
|
|
|
|
|
|
def run_parameter_sweep(
|
|
backends: list[str],
|
|
batch_specs: list[str],
|
|
base_config_args: dict,
|
|
sweep: ParameterSweep,
|
|
console: Console,
|
|
) -> list[BenchmarkResult]:
|
|
"""Run parameter sweep for given backends and batch specs.
|
|
|
|
Args:
|
|
backends: List of backend names
|
|
batch_specs: List of batch specifications
|
|
base_config_args: Base configuration arguments (num_layers, head_dim, etc.)
|
|
sweep: ParameterSweep configuration
|
|
console: Rich console for output
|
|
|
|
Returns:
|
|
List of BenchmarkResult objects
|
|
|
|
"""
|
|
all_results = []
|
|
|
|
# Build list of values to sweep (including auto if requested)
|
|
sweep_values = list(sweep.values)
|
|
if sweep.include_auto:
|
|
sweep_values.append("auto")
|
|
|
|
console.print(f"[yellow]Sweep mode: testing {sweep.param_name} = {sweep_values}[/]")
|
|
|
|
total = len(backends) * len(batch_specs) * len(sweep_values)
|
|
|
|
with tqdm(total=total, desc="Benchmarking") as pbar:
|
|
for backend in backends:
|
|
for spec in batch_specs:
|
|
for value in sweep_values:
|
|
# Create config with original backend for running
|
|
config = BenchmarkConfig(
|
|
backend=backend, batch_spec=spec, **base_config_args
|
|
)
|
|
|
|
# Prepare kwargs for benchmark runner
|
|
kwargs = {}
|
|
if value != "auto":
|
|
kwargs[sweep.param_name] = value
|
|
|
|
# Run benchmark
|
|
result = run_benchmark(config, **kwargs)
|
|
|
|
# Replace backend with labeled version for display
|
|
backend_label = sweep.get_label(backend, value)
|
|
labeled_config = replace(result.config, backend=backend_label)
|
|
result = replace(result, config=labeled_config)
|
|
all_results.append(result)
|
|
|
|
if not result.success:
|
|
console.print(
|
|
f"[red]Error {backend} {spec} {sweep.param_name}="
|
|
f"{value}: {result.error}[/]"
|
|
)
|
|
|
|
pbar.update(1)
|
|
|
|
if base_config_args.get("ncu_profile"):
|
|
return all_results
|
|
|
|
# Display sweep results
|
|
console.print("\n[bold green]Sweep Results:[/]")
|
|
backend_labels = [sweep.get_label(b, v) for b in backends for v in sweep_values]
|
|
formatter = ResultsFormatter(console)
|
|
formatter.print_table(all_results, backend_labels)
|
|
|
|
# Show optimal values
|
|
console.print(f"\n[bold cyan]Optimal {sweep.param_name} per batch spec:[/]")
|
|
by_spec = {}
|
|
for r in all_results:
|
|
if r.success:
|
|
spec = r.config.batch_spec
|
|
if spec not in by_spec:
|
|
by_spec[spec] = []
|
|
by_spec[spec].append(r)
|
|
|
|
for spec in sorted(by_spec.keys(), key=batch_spec_sort_key):
|
|
results = by_spec[spec]
|
|
best = min(results, key=lambda r: r.mean_time)
|
|
console.print(
|
|
f" {spec}: [bold green]{best.config.backend}[/] ({best.mean_time:.6f}s)"
|
|
)
|
|
|
|
return all_results
|
|
|
|
|
|
def load_config_from_yaml(config_path: str) -> dict:
|
|
"""Load configuration from YAML file."""
|
|
with open(config_path) as f:
|
|
return yaml.safe_load(f)
|
|
|
|
|
|
def generate_batch_specs_from_ranges(ranges: list[dict]) -> list[str]:
|
|
"""Generate batch specs from range specifications.
|
|
|
|
Args:
|
|
ranges: List of range specifications, each containing:
|
|
- template: Batch spec template (e.g., "q{q_len}kv1k")
|
|
- q_len: Dict with start, stop, step, end_inclusive (optional)
|
|
- Other parameters can also be ranges
|
|
|
|
Returns:
|
|
List of generated batch spec strings
|
|
|
|
Example:
|
|
ranges = [
|
|
{
|
|
"template": "q{q_len}kv1k",
|
|
"q_len": {
|
|
"start": 1,
|
|
"stop": 16,
|
|
"step": 1,
|
|
"end_inclusive": true # Optional, defaults to true
|
|
}
|
|
}
|
|
]
|
|
Returns: ["q1kv1k", "q2kv1k", ..., "q16kv1k"]
|
|
|
|
"""
|
|
all_specs = []
|
|
|
|
for range_spec in ranges:
|
|
template = range_spec.get("template")
|
|
if not template:
|
|
raise ValueError("Range specification must include 'template'")
|
|
|
|
# Extract all range parameters from the spec
|
|
range_params = {}
|
|
for key, value in range_spec.items():
|
|
if key == "template":
|
|
continue
|
|
if isinstance(value, dict) and "start" in value:
|
|
# This is a range specification
|
|
start = value["start"]
|
|
stop = value["stop"]
|
|
step = value.get("step", 1)
|
|
# Check if end should be inclusive (default: True)
|
|
end_inclusive = value.get("end_inclusive", True)
|
|
|
|
# Adjust stop based on end_inclusive
|
|
if end_inclusive:
|
|
range_params[key] = list(range(start, stop + 1, step))
|
|
else:
|
|
range_params[key] = list(range(start, stop, step))
|
|
else:
|
|
# This is a fixed value
|
|
range_params[key] = [value]
|
|
|
|
# Generate all combinations (Cartesian product)
|
|
if range_params:
|
|
import itertools
|
|
|
|
param_names = list(range_params.keys())
|
|
param_values = [range_params[name] for name in param_names]
|
|
|
|
for values in itertools.product(*param_values):
|
|
params = dict(zip(param_names, values))
|
|
spec = template.format(**params)
|
|
all_specs.append(spec)
|
|
else:
|
|
# No parameters, just use template as-is
|
|
all_specs.append(template)
|
|
|
|
return all_specs
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Universal vLLM attention benchmark",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog=__doc__,
|
|
)
|
|
|
|
# Config file
|
|
parser.add_argument(
|
|
"--config",
|
|
help="Path to YAML config file (overrides other args)",
|
|
)
|
|
|
|
# Backend selection
|
|
parser.add_argument(
|
|
"--backends",
|
|
"--decode-backends",
|
|
nargs="+",
|
|
help="Decode backends to benchmark (flash, triton, flashinfer, cutlass_mla, "
|
|
"flashinfer_mla, flashattn_mla, flashmla)",
|
|
)
|
|
parser.add_argument(
|
|
"--backend",
|
|
help="Single backend (alternative to --backends)",
|
|
)
|
|
parser.add_argument(
|
|
"--prefill-backends",
|
|
nargs="+",
|
|
help="Prefill backends to compare (fa2, fa3, fa4). "
|
|
"Uses the first decode backend for impl construction.",
|
|
)
|
|
parser.add_argument(
|
|
"--fp8-output-scale",
|
|
type=float,
|
|
help="Static per-tensor scale enabling the MLA prefill FP8-output "
|
|
"comparison on FA4 (fused write vs standalone post-quant).",
|
|
)
|
|
parser.add_argument(
|
|
"--fuse-quant-op",
|
|
nargs="+",
|
|
type=_str2bool,
|
|
help="FP8-output write path(s) to run: false = bf16 attention + "
|
|
"standalone static-FP8 quant, true = FA4 writes FP8 directly. "
|
|
"Default: both.",
|
|
)
|
|
|
|
# Batch specifications
|
|
parser.add_argument(
|
|
"--batch-specs",
|
|
nargs="+",
|
|
default=None,
|
|
help="Batch specifications using extended grammar",
|
|
)
|
|
|
|
# Model config
|
|
parser.add_argument("--num-layers", type=int, default=10, help="Number of layers")
|
|
parser.add_argument("--head-dim", type=int, default=128, help="Head dimension")
|
|
parser.add_argument("--num-q-heads", type=int, default=32, help="Query heads")
|
|
parser.add_argument("--num-kv-heads", type=int, default=8, help="KV heads")
|
|
parser.add_argument("--block-size", type=int, default=16, help="Block size")
|
|
parser.add_argument(
|
|
"--v-head-dim",
|
|
type=int,
|
|
default=None,
|
|
help="Value head dimension (defaults to --head-dim if unset)",
|
|
)
|
|
|
|
# MLA-specific model dimensions
|
|
parser.add_argument(
|
|
"--kv-lora-rank", type=int, default=None, help="MLA KV LoRA rank"
|
|
)
|
|
parser.add_argument(
|
|
"--qk-nope-head-dim", type=int, default=None, help="MLA non-RoPE QK head dim"
|
|
)
|
|
parser.add_argument(
|
|
"--qk-rope-head-dim", type=int, default=None, help="MLA RoPE QK head dim"
|
|
)
|
|
|
|
# Benchmark settings
|
|
parser.add_argument("--device", default="cuda:0", help="Device")
|
|
parser.add_argument(
|
|
"--model",
|
|
default="meta-llama/Meta-Llama-3-8B",
|
|
help="Model whose config the backend is built with",
|
|
)
|
|
parser.add_argument(
|
|
"--warmup-ms",
|
|
type=int,
|
|
default=None,
|
|
help=(
|
|
"Warmup window in ms for triton's do_bench (default: triton's own). "
|
|
"Has no effect with CUDA graphs; pass --no-cuda-graphs to use it."
|
|
),
|
|
)
|
|
parser.add_argument("--profile-memory", action="store_true", help="Profile memory")
|
|
parser.add_argument(
|
|
"--kv-cache-dtype",
|
|
default="auto",
|
|
choices=get_args(CacheDType),
|
|
help="KV cache dtype",
|
|
)
|
|
parser.add_argument(
|
|
"--cuda-graphs",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=True,
|
|
help=(
|
|
"Use triton do_bench_cudagraph (True) or do_bench (False) "
|
|
"for timing. CUDA graphs eliminate CPU launch overhead "
|
|
"(default: True)"
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--num-splits",
|
|
type=int,
|
|
default=None,
|
|
help="FlashAttention split-K factor (0=auto heuristic, 1=disabled, >1=force N)",
|
|
)
|
|
parser.add_argument(
|
|
"--ncu-profile",
|
|
action="store_true",
|
|
default=False,
|
|
help=(
|
|
"Enable Nsight Compute profiling mode. Automatically wraps the "
|
|
"script with ncu, capturing a profile with source correlation. "
|
|
"Use --ncu-output to set the output file name."
|
|
),
|
|
)
|
|
parser.add_argument(
|
|
"--ncu-output",
|
|
type=str,
|
|
default="profile",
|
|
help="Output file name for ncu profile (default: 'profile').",
|
|
)
|
|
parser.add_argument(
|
|
"--torch-profile",
|
|
action="store_true",
|
|
default=False,
|
|
help="Collect a PyTorch profiler Chrome trace for each benchmark run.",
|
|
)
|
|
parser.add_argument(
|
|
"--torch-profile-dir",
|
|
default=None,
|
|
help="Directory for PyTorch profiler traces.",
|
|
)
|
|
parser.add_argument(
|
|
"--torch-profile-iters",
|
|
type=int,
|
|
default=3,
|
|
help="Number of forward passes to record per PyTorch profiler trace.",
|
|
)
|
|
parser.add_argument(
|
|
"--sparse-mla-mha-variants",
|
|
nargs="+",
|
|
default=None,
|
|
choices=["dense_mha", "masked_mha", "mqa"],
|
|
help="Sparse MLA variants to run in mha_vs_mqa mode. Defaults to both.",
|
|
)
|
|
|
|
# Parameter sweep (use YAML config for advanced sweeps)
|
|
parser.add_argument(
|
|
"--sweep-param",
|
|
help="Parameter name to sweep (e.g., num_kv_splits, reorder_batch_threshold)",
|
|
)
|
|
parser.add_argument(
|
|
"--sweep-values",
|
|
type=int,
|
|
nargs="+",
|
|
help="Values to sweep for the parameter",
|
|
)
|
|
|
|
# Output
|
|
parser.add_argument("--output-csv", help="Save to CSV")
|
|
parser.add_argument("--output-json", help="Save to JSON")
|
|
|
|
args = parser.parse_args()
|
|
|
|
console = Console()
|
|
console.print("[bold cyan]vLLM Attention Benchmark[/]")
|
|
|
|
# Load config from YAML if provided
|
|
if args.config:
|
|
console.print(f"[yellow]Loading config from: {args.config}[/]")
|
|
yaml_config = load_config_from_yaml(args.config)
|
|
|
|
# Show description if available
|
|
if "description" in yaml_config:
|
|
console.print(f"[dim]{yaml_config['description']}[/]")
|
|
|
|
# Override args with YAML values, but CLI args take precedence
|
|
# Check if CLI provided backends (they would be non-None and not default)
|
|
cli_backends_provided = args.backend is not None or args.backends is not None
|
|
|
|
# Backend(s) - only use YAML if CLI didn't specify
|
|
if not cli_backends_provided:
|
|
if "backend" in yaml_config:
|
|
args.backend = yaml_config["backend"]
|
|
args.backends = None
|
|
elif "backends" in yaml_config:
|
|
args.backends = yaml_config["backends"]
|
|
args.backend = None
|
|
elif "decode_backends" in yaml_config:
|
|
args.backends = yaml_config["decode_backends"]
|
|
args.backend = None
|
|
|
|
# Prefill backends (e.g., ["fa3", "fa4"])
|
|
args.prefill_backends = yaml_config.get("prefill_backends", None)
|
|
args.prefill_backend = yaml_config.get("prefill_backend", None)
|
|
|
|
# FP8 output benchmark knobs; CLI wins.
|
|
if args.fp8_output_scale is None:
|
|
args.fp8_output_scale = yaml_config.get("fp8_output_scale", None)
|
|
if args.fuse_quant_op is None:
|
|
args.fuse_quant_op = yaml_config.get("fuse_quant_op", None)
|
|
|
|
# Check for special modes
|
|
args.mode = yaml_config.get("mode", None)
|
|
|
|
# Batch specs and sizes
|
|
# Support both explicit batch_specs and generated batch_spec_ranges
|
|
# CLI --batch-specs takes precedence over YAML when provided.
|
|
cli_batch_specs_provided = args.batch_specs is not None
|
|
if not cli_batch_specs_provided:
|
|
if "batch_spec_ranges" in yaml_config:
|
|
# Generate batch specs from ranges
|
|
generated_specs = generate_batch_specs_from_ranges(
|
|
yaml_config["batch_spec_ranges"]
|
|
)
|
|
# Combine with any explicit batch_specs
|
|
if "batch_specs" in yaml_config:
|
|
args.batch_specs = yaml_config["batch_specs"] + generated_specs
|
|
else:
|
|
args.batch_specs = generated_specs
|
|
console.print(
|
|
f"[dim]Generated {len(generated_specs)} batch specs from ranges[/]"
|
|
)
|
|
elif "batch_specs" in yaml_config:
|
|
args.batch_specs = yaml_config["batch_specs"]
|
|
|
|
args.batch_sizes = yaml_config.get("batch_sizes", None)
|
|
|
|
# Model config
|
|
if "model" in yaml_config:
|
|
model = yaml_config["model"]
|
|
args.num_layers = model.get("num_layers", args.num_layers)
|
|
args.head_dim = model.get("head_dim", args.head_dim)
|
|
args.v_head_dim = model.get("v_head_dim", args.v_head_dim)
|
|
args.num_q_heads = model.get("num_q_heads", args.num_q_heads)
|
|
args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads)
|
|
args.block_size = model.get("block_size", args.block_size)
|
|
args.max_model_len = model.get(
|
|
"max_model_len", getattr(args, "max_model_len", None)
|
|
)
|
|
# MLA-specific dimensions
|
|
args.kv_lora_rank = model.get("kv_lora_rank", args.kv_lora_rank)
|
|
args.qk_nope_head_dim = model.get("qk_nope_head_dim", args.qk_nope_head_dim)
|
|
args.qk_rope_head_dim = model.get("qk_rope_head_dim", args.qk_rope_head_dim)
|
|
|
|
# Benchmark settings (top-level keys)
|
|
if "device" in yaml_config:
|
|
args.device = yaml_config["device"]
|
|
if "warmup_ms" in yaml_config:
|
|
args.warmup_ms = yaml_config["warmup_ms"]
|
|
if "profile_memory" in yaml_config:
|
|
args.profile_memory = yaml_config["profile_memory"]
|
|
if "kv_cache_dtype" in yaml_config:
|
|
args.kv_cache_dtype = yaml_config["kv_cache_dtype"]
|
|
if "cuda_graphs" in yaml_config:
|
|
args.cuda_graphs = yaml_config["cuda_graphs"]
|
|
if "ncu_profile" in yaml_config:
|
|
args.ncu_profile = yaml_config["ncu_profile"]
|
|
if "torch_profile" in yaml_config:
|
|
args.torch_profile = yaml_config["torch_profile"]
|
|
if "torch_profile_dir" in yaml_config:
|
|
args.torch_profile_dir = yaml_config["torch_profile_dir"]
|
|
if "torch_profile_iters" in yaml_config:
|
|
args.torch_profile_iters = yaml_config["torch_profile_iters"]
|
|
args.sparse_mla_topk_pattern = yaml_config.get(
|
|
"sparse_mla_topk_pattern", "random"
|
|
)
|
|
args.sparse_mla_dense_mha_max_seq_len = yaml_config.get(
|
|
"sparse_mla_dense_mha_max_seq_len", None
|
|
)
|
|
args.sparse_mla_mha_variants = yaml_config.get(
|
|
"sparse_mla_mha_variants", args.sparse_mla_mha_variants
|
|
)
|
|
args.sparse_mla_masked_mha_max_seq_len = yaml_config.get(
|
|
"sparse_mla_masked_mha_max_seq_len", None
|
|
)
|
|
|
|
# Parameter sweep configuration
|
|
if "parameter_sweep" in yaml_config:
|
|
sweep_config = yaml_config["parameter_sweep"]
|
|
args.parameter_sweep = ParameterSweep(
|
|
param_name=sweep_config["param_name"],
|
|
values=sweep_config["values"],
|
|
include_auto=sweep_config.get("include_auto", False),
|
|
label_format=sweep_config.get(
|
|
"label_format", "{backend}_{param_name}_{value}"
|
|
),
|
|
)
|
|
else:
|
|
args.parameter_sweep = None
|
|
|
|
# Model parameter sweep configuration
|
|
if "model_parameter_sweep" in yaml_config:
|
|
sweep_config = yaml_config["model_parameter_sweep"]
|
|
args.model_parameter_sweep = ModelParameterSweep(
|
|
param_name=sweep_config.get("param_name"),
|
|
values=sweep_config["values"],
|
|
label_format=sweep_config.get(
|
|
"label_format", "{backend}_{param_name}_{value}"
|
|
),
|
|
)
|
|
else:
|
|
args.model_parameter_sweep = None
|
|
|
|
# Output
|
|
if "output" in yaml_config:
|
|
output = yaml_config["output"]
|
|
if "csv" in output and not args.output_csv:
|
|
args.output_csv = output["csv"]
|
|
if "json" in output and not args.output_json:
|
|
args.output_json = output["json"]
|
|
|
|
console.print()
|
|
|
|
# Re-exec under ncu if --ncu-profile and not already inside ncu. This runs
|
|
# after YAML processing so ncu_profile set via config file is honored.
|
|
if args.ncu_profile and "_NCU_INNER" not in os.environ:
|
|
ncu = shutil.which("ncu")
|
|
if ncu is None:
|
|
print("Error: 'ncu' not found in PATH", file=sys.stderr)
|
|
sys.exit(1)
|
|
cmd = [
|
|
ncu,
|
|
"--profile-from-start",
|
|
"off",
|
|
"--set",
|
|
"full",
|
|
"--import-source",
|
|
"yes",
|
|
"-o",
|
|
args.ncu_output,
|
|
sys.executable,
|
|
*sys.argv,
|
|
]
|
|
env = os.environ.copy()
|
|
env["CUTE_DSL_LINEINFO"] = "1"
|
|
env["_NCU_INNER"] = "1"
|
|
print(f"Launching: {' '.join(cmd)}")
|
|
sys.exit(subprocess.call(cmd, env=env))
|
|
|
|
# Handle CLI-based parameter sweep (if not from YAML)
|
|
if (
|
|
(not hasattr(args, "parameter_sweep") or args.parameter_sweep is None)
|
|
and args.sweep_param
|
|
and args.sweep_values
|
|
):
|
|
args.parameter_sweep = ParameterSweep(
|
|
param_name=args.sweep_param,
|
|
values=args.sweep_values,
|
|
include_auto=False,
|
|
label_format="{backend}_{param_name}_{value}",
|
|
)
|
|
|
|
# Determine backends
|
|
backends = args.backends or ([args.backend] if args.backend else ["flash"])
|
|
prefill_backends = getattr(args, "prefill_backends", None)
|
|
if not args.batch_specs:
|
|
args.batch_specs = ["q2k", "8q1s1k"]
|
|
console.print(f"Backends: {', '.join(backends)}")
|
|
if prefill_backends:
|
|
console.print(f"Prefill backends: {', '.join(prefill_backends)}")
|
|
console.print(f"Batch specs: {', '.join(args.batch_specs)}")
|
|
console.print(f"KV cache dtype: {args.kv_cache_dtype}")
|
|
console.print(f"CUDA graphs: {args.cuda_graphs}")
|
|
if args.warmup_ms is not None and args.cuda_graphs:
|
|
console.print(
|
|
"[yellow]Warning: --warmup-ms is ignored with CUDA graphs "
|
|
"(do_bench_cudagraph warms up internally). Pass --no-cuda-graphs "
|
|
"to use it.[/]"
|
|
)
|
|
if args.num_splits == 0 and args.cuda_graphs:
|
|
console.print(
|
|
"[yellow]Warning: --num-splits 0 (FA3 heuristic) is not CUDA-graph "
|
|
"compatible and may fail or fall back. Pass --no-cuda-graphs or use "
|
|
"--num-splits >=1.[/]"
|
|
)
|
|
console.print()
|
|
|
|
init_workspace_manager(args.device)
|
|
|
|
# Run benchmarks
|
|
all_results = []
|
|
|
|
# Under ncu profiling the kernels run only to be captured by the profiler;
|
|
# timings are placeholder zeros, so the result tables and saved metrics are
|
|
# skipped. The Nsight Compute report (--ncu-output) holds the real data.
|
|
if args.ncu_profile:
|
|
console.print(
|
|
"[dim]ncu profiling enabled: result tables and saved metrics are "
|
|
"skipped (timings are placeholder zeros).[/]"
|
|
)
|
|
|
|
# FA4 fused FP8 output vs standalone post-quant, on the same fa4 kernel:
|
|
# the delta is the post-quant kernel the fused path removes.
|
|
fp8_output_scale = getattr(args, "fp8_output_scale", None)
|
|
if fp8_output_scale is not None:
|
|
decode_backend = backends[0]
|
|
fuse_variants = args.fuse_quant_op or [False, True]
|
|
label_of = {False: "post_quant", True: "fused"}
|
|
console.print(
|
|
f"[yellow]FP8 output comparison @ scale={fp8_output_scale} "
|
|
f"(prefill=fa4, decode impl={decode_backend})[/]"
|
|
)
|
|
fp8_results = []
|
|
total = len(fuse_variants) * len(args.batch_specs)
|
|
with tqdm(total=total, desc="FP8 output benchmarking") as pbar:
|
|
for spec in args.batch_specs:
|
|
for fuse in fuse_variants:
|
|
config = BenchmarkConfig(
|
|
backend=decode_backend,
|
|
batch_spec=spec,
|
|
num_layers=args.num_layers,
|
|
head_dim=args.head_dim,
|
|
num_q_heads=args.num_q_heads,
|
|
num_kv_heads=args.num_kv_heads,
|
|
block_size=args.block_size,
|
|
device=args.device,
|
|
profile_memory=args.profile_memory,
|
|
kv_cache_dtype=args.kv_cache_dtype,
|
|
model=args.model,
|
|
use_cuda_graphs=args.cuda_graphs,
|
|
prefill_backend="fa4",
|
|
)
|
|
result = run_benchmark(
|
|
config, output_scale=fp8_output_scale, fuse_quant_op=fuse
|
|
)
|
|
label = label_of[fuse]
|
|
labeled_config = replace(result.config, backend=label)
|
|
result = replace(result, config=labeled_config)
|
|
fp8_results.append(result)
|
|
|
|
if not result.success:
|
|
console.print(f"[red]Error {label} {spec}: {result.error}[/]")
|
|
|
|
pbar.update(1)
|
|
|
|
console.print("\n[bold green]FP8 Output Results:[/]")
|
|
formatter = ResultsFormatter(console)
|
|
labels = [label_of[f] for f in fuse_variants]
|
|
formatter.print_table(fp8_results, labels, compare_to_fastest=True)
|
|
all_results = fp8_results
|
|
|
|
# Handle special mode: decode_vs_prefill comparison
|
|
elif hasattr(args, "mode") and args.mode == "decode_vs_prefill":
|
|
console.print("[yellow]Mode: Decode vs Prefill pipeline comparison[/]")
|
|
console.print(
|
|
"[dim]For each query length, testing both decode and prefill pipelines[/]"
|
|
)
|
|
console.print("[dim]Using batched execution for optimal performance[/]")
|
|
|
|
# Extract batch sizes from config
|
|
batch_sizes = getattr(args, "batch_sizes", [1])
|
|
backend = backends[0] # Use first backend (should only be one)
|
|
|
|
# Calculate total benchmarks
|
|
total = len(batch_sizes)
|
|
|
|
with tqdm(total=total, desc="Benchmarking") as pbar:
|
|
for batch_size in batch_sizes:
|
|
# Prepare all configs for this batch size
|
|
configs_with_thresholds = []
|
|
|
|
for spec in args.batch_specs:
|
|
# Parse the batch spec to get query length
|
|
requests = parse_batch_spec(spec)
|
|
if not requests:
|
|
console.print(
|
|
f"[red]Error: Could not parse batch spec '{spec}'[/]"
|
|
)
|
|
continue
|
|
|
|
# Get query length from first request
|
|
query_length = requests[0].q_len
|
|
|
|
# Create batch spec for this batch size
|
|
# For batch_size > 1, we need to prepend the count
|
|
batch_spec = f"{batch_size}{spec}" if batch_size > 1 else spec
|
|
|
|
# Create base config (without backend name)
|
|
base_config = BenchmarkConfig(
|
|
backend=backend, # Will be overridden later
|
|
batch_spec=batch_spec,
|
|
num_layers=args.num_layers,
|
|
head_dim=args.head_dim,
|
|
num_q_heads=args.num_q_heads,
|
|
num_kv_heads=args.num_kv_heads,
|
|
block_size=args.block_size,
|
|
device=args.device,
|
|
profile_memory=args.profile_memory,
|
|
kv_cache_dtype=args.kv_cache_dtype,
|
|
model=args.model,
|
|
use_cuda_graphs=args.cuda_graphs,
|
|
ncu_profile=args.ncu_profile,
|
|
warmup_ms=args.warmup_ms,
|
|
)
|
|
|
|
# Add decode pipeline config
|
|
decode_threshold = query_length
|
|
config_decode = replace(
|
|
base_config,
|
|
backend=f"{backend}_decode_qlen{query_length}_bs{batch_size}",
|
|
)
|
|
configs_with_thresholds.append((config_decode, decode_threshold))
|
|
|
|
# Add prefill pipeline config if query_length > 1
|
|
if query_length > 1:
|
|
prefill_threshold = query_length - 1
|
|
config_prefill = replace(
|
|
base_config,
|
|
backend=f"{backend}_prefill_qlen{query_length}"
|
|
f"_bs{batch_size}",
|
|
)
|
|
configs_with_thresholds.append(
|
|
(config_prefill, prefill_threshold)
|
|
)
|
|
|
|
# Run all benchmarks for this batch size in one go (batched mode)
|
|
try:
|
|
from mla_runner import run_mla_benchmark as run_mla
|
|
|
|
# Use batched API: pass list of (config, threshold) tuples
|
|
timing_results = run_mla(backend, configs_with_thresholds)
|
|
|
|
# Create BenchmarkResult objects from timing results
|
|
for (config, _), timing in zip(
|
|
configs_with_thresholds, timing_results
|
|
):
|
|
result = BenchmarkResult(
|
|
config=config,
|
|
mean_time=timing["mean"],
|
|
median_time=timing.get("median", timing["mean"]),
|
|
std_time=timing["std"],
|
|
min_time=timing["min"],
|
|
max_time=timing["max"],
|
|
throughput_tokens_per_sec=timing.get("throughput", None),
|
|
)
|
|
all_results.append(result)
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
console.print(
|
|
f"[red]Error running batched benchmarks for "
|
|
f"batch_size={batch_size}: {e}[/]"
|
|
)
|
|
console.print("[red]Traceback:[/]")
|
|
traceback.print_exc()
|
|
# Add error results for all configs
|
|
for config, _ in configs_with_thresholds:
|
|
result = BenchmarkResult(
|
|
config=config,
|
|
mean_time=float("inf"),
|
|
median_time=float("inf"),
|
|
std_time=0,
|
|
min_time=float("inf"),
|
|
max_time=float("inf"),
|
|
error=str(e),
|
|
)
|
|
all_results.append(result)
|
|
|
|
pbar.update(1)
|
|
|
|
if args.ncu_profile:
|
|
if any(not result.success for result in all_results):
|
|
raise SystemExit(1)
|
|
return
|
|
|
|
# Display decode vs prefill results
|
|
console.print("\n[bold green]Decode vs Prefill Results:[/]")
|
|
|
|
# Group by batch size
|
|
by_batch_size = {}
|
|
for r in all_results:
|
|
if r.success:
|
|
# Extract batch size from backend name
|
|
parts = r.config.backend.split("_")
|
|
bs_part = [p for p in parts if p.startswith("bs")]
|
|
if bs_part:
|
|
bs = int(bs_part[0][2:])
|
|
if bs not in by_batch_size:
|
|
by_batch_size[bs] = []
|
|
by_batch_size[bs].append(r)
|
|
|
|
# For each batch size, analyze crossover point
|
|
for bs in sorted(by_batch_size.keys()):
|
|
console.print(f"\n[bold cyan]Batch size: {bs}[/]")
|
|
results = by_batch_size[bs]
|
|
|
|
# Group by query length
|
|
by_qlen = {}
|
|
for r in results:
|
|
parts = r.config.backend.split("_")
|
|
qlen_part = [p for p in parts if p.startswith("qlen")]
|
|
if qlen_part:
|
|
qlen = int(qlen_part[0][4:])
|
|
if qlen not in by_qlen:
|
|
by_qlen[qlen] = {}
|
|
|
|
pipeline = "decode" if "decode" in r.config.backend else "prefill"
|
|
by_qlen[qlen][pipeline] = r
|
|
|
|
# Find crossover point
|
|
last_decode_faster = None
|
|
for qlen in sorted(by_qlen.keys()):
|
|
pipelines = by_qlen[qlen]
|
|
if "decode" in pipelines and "prefill" in pipelines:
|
|
decode_time = pipelines["decode"].mean_time
|
|
prefill_time = pipelines["prefill"].mean_time
|
|
faster = "decode" if decode_time < prefill_time else "prefill"
|
|
|
|
speedup = (
|
|
prefill_time / decode_time
|
|
if decode_time < prefill_time
|
|
else decode_time / prefill_time
|
|
)
|
|
|
|
console.print(
|
|
f" qlen={qlen:3d}: decode={decode_time:.6f}s, "
|
|
f"prefill={prefill_time:.6f}s -> "
|
|
f"[bold]{faster}[/] ({speedup:.2f}x)"
|
|
)
|
|
|
|
if faster == "decode":
|
|
last_decode_faster = qlen
|
|
|
|
if last_decode_faster is not None:
|
|
optimal_threshold = last_decode_faster
|
|
console.print(
|
|
f"\n [bold green]Optimal threshold for batch_size={bs}: "
|
|
f"{optimal_threshold}[/]"
|
|
)
|
|
console.print(
|
|
f" [dim](Use decode pipeline for query_length <= "
|
|
f"{optimal_threshold})[/]"
|
|
)
|
|
else:
|
|
console.print(
|
|
f"\n [yellow]Prefill always faster for batch_size={bs}[/]"
|
|
)
|
|
|
|
# Handle MHA vs MQA comparison mode for sparse MLA
|
|
elif hasattr(args, "mode") and args.mode == "mha_vs_mqa":
|
|
console.print("[yellow]Mode: MHA vs MQA comparison for sparse MLA[/]")
|
|
|
|
sparse_mla_topk_pattern = getattr(args, "sparse_mla_topk_pattern", "random")
|
|
dense_mha_max_seq_len = getattr(args, "sparse_mla_dense_mha_max_seq_len", None)
|
|
prefill_backend = getattr(args, "prefill_backend", None)
|
|
if prefill_backend:
|
|
console.print(f"Prefill backend: {prefill_backend}")
|
|
available_variants = [
|
|
("dense_mha", False, "dense"),
|
|
("masked_mha", False, "masked"),
|
|
("mqa", True, "auto"),
|
|
]
|
|
requested_variants = getattr(args, "sparse_mla_mha_variants", None)
|
|
if requested_variants is not None:
|
|
valid_variants = {label for label, _, _ in available_variants}
|
|
invalid_variants = sorted(set(requested_variants) - valid_variants)
|
|
if invalid_variants:
|
|
raise ValueError(
|
|
"Invalid sparse_mla_mha_variants entries: "
|
|
f"{invalid_variants}. Valid variants are: "
|
|
f"{sorted(valid_variants)}"
|
|
)
|
|
requested_variant_set = set(requested_variants)
|
|
variants = [
|
|
variant
|
|
for variant in available_variants
|
|
if variant[0] in requested_variant_set
|
|
]
|
|
else:
|
|
variants = available_variants
|
|
model_sweep = getattr(args, "model_parameter_sweep", None)
|
|
if model_sweep is None:
|
|
num_q_heads_values = [args.num_q_heads]
|
|
elif model_sweep.param_name == "num_q_heads":
|
|
num_q_heads_values = model_sweep.values
|
|
else:
|
|
raise ValueError(
|
|
"mha_vs_mqa mode only supports a num_q_heads model parameter sweep"
|
|
)
|
|
benchmark_specs = [
|
|
(spec, num_q_heads)
|
|
for spec in args.batch_specs
|
|
for num_q_heads in num_q_heads_values
|
|
]
|
|
masked_mha_max_seq_len = getattr(
|
|
args, "sparse_mla_masked_mha_max_seq_len", None
|
|
)
|
|
formatter = ResultsFormatter(console)
|
|
total = 0
|
|
for spec, _ in benchmark_specs:
|
|
q_len = max(request.q_len for request in parse_batch_spec(spec))
|
|
for variant_label, _, _ in variants:
|
|
if (
|
|
variant_label == "dense_mha"
|
|
and dense_mha_max_seq_len is not None
|
|
and q_len > dense_mha_max_seq_len
|
|
) or (
|
|
variant_label == "masked_mha"
|
|
and masked_mha_max_seq_len is not None
|
|
and q_len > masked_mha_max_seq_len
|
|
):
|
|
continue
|
|
total += len(backends)
|
|
|
|
with tqdm(total=total, desc="Benchmarking") as pbar:
|
|
for spec, num_q_heads in benchmark_specs:
|
|
q_len = max(request.q_len for request in parse_batch_spec(spec))
|
|
for backend in backends:
|
|
for variant_label, force_mqa, mha_mode in variants:
|
|
if (
|
|
variant_label == "dense_mha"
|
|
and dense_mha_max_seq_len is not None
|
|
and q_len > dense_mha_max_seq_len
|
|
) or (
|
|
variant_label == "masked_mha"
|
|
and masked_mha_max_seq_len is not None
|
|
and q_len > masked_mha_max_seq_len
|
|
):
|
|
continue
|
|
result_backend = f"{backend}_{variant_label}"
|
|
if model_sweep is not None:
|
|
result_backend = model_sweep.get_label(
|
|
result_backend, num_q_heads
|
|
)
|
|
config = BenchmarkConfig(
|
|
backend=result_backend,
|
|
batch_spec=spec,
|
|
num_layers=args.num_layers,
|
|
head_dim=args.head_dim,
|
|
num_q_heads=num_q_heads,
|
|
num_kv_heads=args.num_kv_heads,
|
|
block_size=args.block_size,
|
|
device=args.device,
|
|
max_model_len=getattr(args, "max_model_len", None),
|
|
kv_cache_dtype=args.kv_cache_dtype,
|
|
model=args.model,
|
|
profile_memory=args.profile_memory,
|
|
use_cuda_graphs=args.cuda_graphs,
|
|
ncu_profile=args.ncu_profile,
|
|
torch_profile=args.torch_profile,
|
|
torch_profile_dir=args.torch_profile_dir,
|
|
torch_profile_iters=args.torch_profile_iters,
|
|
warmup_ms=args.warmup_ms,
|
|
kv_lora_rank=getattr(args, "kv_lora_rank", None),
|
|
qk_nope_head_dim=getattr(args, "qk_nope_head_dim", None),
|
|
qk_rope_head_dim=getattr(args, "qk_rope_head_dim", None),
|
|
v_head_dim=getattr(args, "v_head_dim", None),
|
|
sparse_mla_force_mqa=force_mqa,
|
|
sparse_mla_mha_mode=mha_mode,
|
|
sparse_mla_dense_mha_max_seq_len=dense_mha_max_seq_len,
|
|
sparse_mla_topk_pattern=sparse_mla_topk_pattern,
|
|
prefill_backend=prefill_backend,
|
|
sparse_mla_masked_mha_max_seq_len=masked_mha_max_seq_len,
|
|
)
|
|
|
|
# run_mla_benchmark needs the real backend name
|
|
from mla_runner import run_mla_benchmark as run_mla
|
|
|
|
run_label = f"{backend}_{variant_label} {num_q_heads}h {spec}"
|
|
pbar.set_postfix_str(run_label)
|
|
|
|
try:
|
|
result = run_mla(
|
|
backend,
|
|
config,
|
|
prefill_backend=prefill_backend,
|
|
sparse_mla_force_mqa=force_mqa,
|
|
)
|
|
except Exception as e:
|
|
result = BenchmarkResult(
|
|
config=config,
|
|
mean_time=float("inf"),
|
|
median_time=float("inf"),
|
|
std_time=0,
|
|
min_time=float("inf"),
|
|
max_time=float("inf"),
|
|
error=str(e),
|
|
)
|
|
|
|
all_results.append(result)
|
|
if args.output_csv:
|
|
formatter.save_csv(all_results, args.output_csv)
|
|
if args.output_json:
|
|
formatter.save_json(all_results, args.output_json)
|
|
|
|
if not result.success:
|
|
console.print(
|
|
f"[red]Error {backend}_{variant_label} "
|
|
f"{spec}: {result.error}[/]"
|
|
)
|
|
|
|
pbar.update(1)
|
|
|
|
# Display results with variant labels as separate "backends"
|
|
console.print("\n[bold green]MHA vs MQA Results:[/]")
|
|
variant_backends = []
|
|
for num_q_heads in num_q_heads_values:
|
|
for backend in backends:
|
|
for variant, _, _ in variants:
|
|
label = f"{backend}_{variant}"
|
|
if model_sweep is not None:
|
|
label = model_sweep.get_label(label, num_q_heads)
|
|
variant_backends.append(label)
|
|
formatter.print_table(all_results, variant_backends)
|
|
|
|
# Handle model parameter sweep mode
|
|
elif hasattr(args, "model_parameter_sweep") and args.model_parameter_sweep:
|
|
# Model parameter sweep
|
|
base_config_args = {
|
|
"num_layers": args.num_layers,
|
|
"head_dim": args.head_dim,
|
|
"v_head_dim": args.v_head_dim,
|
|
"num_q_heads": args.num_q_heads,
|
|
"num_kv_heads": args.num_kv_heads,
|
|
"block_size": args.block_size,
|
|
"device": args.device,
|
|
"profile_memory": args.profile_memory,
|
|
"kv_cache_dtype": args.kv_cache_dtype,
|
|
"use_cuda_graphs": args.cuda_graphs,
|
|
"ncu_profile": args.ncu_profile,
|
|
"warmup_ms": args.warmup_ms,
|
|
"num_splits": args.num_splits,
|
|
"kv_lora_rank": args.kv_lora_rank,
|
|
"qk_nope_head_dim": args.qk_nope_head_dim,
|
|
"qk_rope_head_dim": args.qk_rope_head_dim,
|
|
}
|
|
all_results = run_model_parameter_sweep(
|
|
backends,
|
|
args.batch_specs,
|
|
base_config_args,
|
|
args.model_parameter_sweep,
|
|
console,
|
|
)
|
|
|
|
# Handle parameter sweep mode (unified)
|
|
elif hasattr(args, "parameter_sweep") and args.parameter_sweep:
|
|
# Unified parameter sweep
|
|
base_config_args = {
|
|
"num_layers": args.num_layers,
|
|
"head_dim": args.head_dim,
|
|
"v_head_dim": args.v_head_dim,
|
|
"num_q_heads": args.num_q_heads,
|
|
"num_kv_heads": args.num_kv_heads,
|
|
"block_size": args.block_size,
|
|
"device": args.device,
|
|
"profile_memory": args.profile_memory,
|
|
"kv_cache_dtype": args.kv_cache_dtype,
|
|
"use_cuda_graphs": args.cuda_graphs,
|
|
"ncu_profile": args.ncu_profile,
|
|
"warmup_ms": args.warmup_ms,
|
|
"num_splits": args.num_splits,
|
|
}
|
|
all_results = run_parameter_sweep(
|
|
backends, args.batch_specs, base_config_args, args.parameter_sweep, console
|
|
)
|
|
|
|
else:
|
|
# Normal mode: compare backends
|
|
decode_results = []
|
|
prefill_results = []
|
|
|
|
# Run decode backend comparison
|
|
if not prefill_backends:
|
|
# No prefill backends specified: compare decode backends as before
|
|
total = len(backends) * len(args.batch_specs)
|
|
|
|
with tqdm(total=total, desc="Benchmarking") as pbar:
|
|
for spec in args.batch_specs:
|
|
for backend in backends:
|
|
config = BenchmarkConfig(
|
|
backend=backend,
|
|
batch_spec=spec,
|
|
num_layers=args.num_layers,
|
|
head_dim=args.head_dim,
|
|
v_head_dim=getattr(args, "v_head_dim", None),
|
|
num_q_heads=args.num_q_heads,
|
|
num_kv_heads=args.num_kv_heads,
|
|
block_size=args.block_size,
|
|
device=args.device,
|
|
profile_memory=args.profile_memory,
|
|
kv_cache_dtype=args.kv_cache_dtype,
|
|
model=args.model,
|
|
use_cuda_graphs=args.cuda_graphs,
|
|
ncu_profile=args.ncu_profile,
|
|
warmup_ms=args.warmup_ms,
|
|
num_splits=args.num_splits,
|
|
)
|
|
|
|
result = run_benchmark(config)
|
|
decode_results.append(result)
|
|
|
|
if not result.success:
|
|
console.print(
|
|
f"[red]Error {backend} {spec}: {result.error}[/]"
|
|
)
|
|
|
|
pbar.update(1)
|
|
|
|
if not args.ncu_profile:
|
|
console.print("\n[bold green]Results:[/]")
|
|
formatter = ResultsFormatter(console)
|
|
formatter.print_table(decode_results, backends)
|
|
|
|
# Run prefill backend comparison
|
|
if prefill_backends:
|
|
# Use first decode backend for impl construction
|
|
decode_backend = backends[0]
|
|
total = len(prefill_backends) * len(args.batch_specs)
|
|
|
|
console.print(
|
|
f"[yellow]Prefill comparison mode: "
|
|
f"using {decode_backend} for decode impl[/]"
|
|
)
|
|
|
|
with tqdm(total=total, desc="Prefill benchmarking") as pbar:
|
|
for spec in args.batch_specs:
|
|
for pb in prefill_backends:
|
|
config = BenchmarkConfig(
|
|
backend=decode_backend,
|
|
batch_spec=spec,
|
|
num_layers=args.num_layers,
|
|
head_dim=args.head_dim,
|
|
num_q_heads=args.num_q_heads,
|
|
num_kv_heads=args.num_kv_heads,
|
|
block_size=args.block_size,
|
|
device=args.device,
|
|
profile_memory=args.profile_memory,
|
|
warmup_ms=args.warmup_ms,
|
|
prefill_backend=pb,
|
|
kv_lora_rank=args.kv_lora_rank,
|
|
qk_nope_head_dim=args.qk_nope_head_dim,
|
|
qk_rope_head_dim=args.qk_rope_head_dim,
|
|
v_head_dim=args.v_head_dim,
|
|
)
|
|
|
|
result = run_benchmark(config)
|
|
|
|
# Label result with prefill backend name for display
|
|
labeled_config = replace(result.config, backend=pb)
|
|
result = replace(result, config=labeled_config)
|
|
prefill_results.append(result)
|
|
|
|
if not result.success:
|
|
console.print(f"[red]Error {pb} {spec}: {result.error}[/]")
|
|
|
|
pbar.update(1)
|
|
|
|
if not args.ncu_profile:
|
|
console.print("\n[bold green]Prefill Backend Results:[/]")
|
|
formatter = ResultsFormatter(console)
|
|
formatter.print_table(
|
|
prefill_results, prefill_backends, compare_to_fastest=True
|
|
)
|
|
|
|
all_results = decode_results + prefill_results
|
|
|
|
# Save results (skip ncu profiling runs: timings are placeholder zeros)
|
|
if all_results and not args.ncu_profile:
|
|
formatter = ResultsFormatter(console)
|
|
if args.output_csv:
|
|
formatter.save_csv(all_results, args.output_csv)
|
|
if args.output_json:
|
|
formatter.save_json(all_results, args.output_json)
|
|
|
|
if any(not result.success for result in all_results):
|
|
raise SystemExit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|