run_benchmark.py

#!/usr/bin/env python3
"""
CLI entry point for HRM SRAM/DRAM Memory Tiering Benchmarks.

Usage:
    # Compare tiered vs baseline HRM
    python run_benchmark.py --mode compare --batch-sizes 1,8,32 --seq-lens 64,128

    # Benchmark tiered model only
    python run_benchmark.py --mode tiered --warmup 5 --iterations 50 --output results.json

    # Generate plots
    python run_benchmark.py --mode compare --plot --output-dir benchmark_results/

    # Quick smoke test
    python run_benchmark.py --mode tiered --warmup 1 --iterations 3 --batch-sizes 2 --seq-lens 16
"""

import argparse
import json
import os
import sys
from dataclasses import asdict

# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from benchmark import (
    benchmark_tiered_model,
    benchmark_baseline_model,
    compare_models,
    print_results_table,
    generate_plots,
    BenchmarkResult,
)


def parse_int_list(s: str):
    return [int(x.strip()) for x in s.split(',')]


def main():
    parser = argparse.ArgumentParser(
        description='HRM SRAM/DRAM Memory Tiering Benchmark Suite (Triton)',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )

    parser.add_argument(
        '--mode', choices=['tiered', 'baseline', 'compare'],
        default='compare',
        help='Benchmark mode: tiered-only, baseline-only, or comparison (default: compare)',
    )
    parser.add_argument(
        '--batch-sizes', type=str, default='1,8,32',
        help='Comma-separated batch sizes to benchmark (default: 1,8,32)',
    )
    parser.add_argument(
        '--seq-lens', type=str, default='64,128',
        help='Comma-separated sequence lengths (default: 64,128)',
    )
    parser.add_argument(
        '--hidden-size', type=int, default=512,
        help='Model hidden size (default: 512)',
    )
    parser.add_argument(
        '--num-heads', type=int, default=8,
        help='Number of attention heads (default: 8)',
    )
    parser.add_argument(
        '--H-cycles', type=int, default=2,
        help='H-level recurrence cycles (default: 2)',
    )
    parser.add_argument(
        '--L-cycles', type=int, default=2,
        help='L-level recurrence cycles (default: 2)',
    )
    parser.add_argument(
        '--H-layers', type=int, default=4,
        help='H-level transformer layers (default: 4)',
    )
    parser.add_argument(
        '--L-layers', type=int, default=4,
        help='L-level transformer layers (default: 4)',
    )
    parser.add_argument(
        '--warmup', type=int, default=5,
        help='Warmup iterations (default: 5)',
    )
    parser.add_argument(
        '--iterations', type=int, default=20,
        help='Benchmark iterations (default: 20)',
    )
    parser.add_argument(
        '--output', type=str, default=None,
        help='Output JSON file path for results',
    )
    parser.add_argument(
        '--output-dir', type=str, default='benchmark_results',
        help='Directory for output files and plots (default: benchmark_results/)',
    )
    parser.add_argument(
        '--plot', action='store_true',
        help='Generate comparison plots (requires matplotlib)',
    )
    parser.add_argument(
        '--device', type=str, default=None,
        help='Device: cuda, cpu, or cuda:N (default: auto-detect)',
    )

    args = parser.parse_args()

    batch_sizes = parse_int_list(args.batch_sizes)
    seq_lens = parse_int_list(args.seq_lens)

    # Device setup
    import torch
    if args.device:
        device = torch.device(args.device)
    else:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    print(f"\n{'='*60}")
    print(f"  HRM SRAM/DRAM Memory Tiering Benchmark")
    print(f"  Device: {device}")
    print(f"  Mode: {args.mode}")
    print(f"  Batch sizes: {batch_sizes}")
    print(f"  Sequence lengths: {seq_lens}")
    print(f"  Hidden size: {args.hidden_size}")
    print(f"  H/L cycles: {args.H_cycles}/{args.L_cycles}")
    print(f"  H/L layers: {args.H_layers}/{args.L_layers}")
    print(f"  Warmup: {args.warmup}, Iterations: {args.iterations}")
    print(f"{'='*60}")

    if args.mode == 'compare':
        results = compare_models(
            batch_sizes=batch_sizes,
            seq_lens=seq_lens,
            hidden_size=args.hidden_size,
            warmup=args.warmup,
            iterations=args.iterations,
            device=device,
        )

        if args.plot:
            generate_plots(results, output_dir=args.output_dir)

        # Save results
        output_path = args.output or os.path.join(args.output_dir, 'results.json')
        os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)
        with open(output_path, 'w') as f:
            json.dump(results, f, indent=2, default=str)
        print(f"\n  Results saved: {output_path}")

    elif args.mode == 'tiered':
        all_results = []
        for bs in batch_sizes:
            for sl in seq_lens:
                print(f"\n  Benchmarking tiered model: bs={bs}, seq={sl}")
                r = benchmark_tiered_model(
                    batch_size=bs, seq_len=sl,
                    hidden_size=args.hidden_size,
                    num_heads=args.num_heads,
                    H_cycles=args.H_cycles, L_cycles=args.L_cycles,
                    H_layers=args.H_layers, L_layers=args.L_layers,
                    warmup=args.warmup, iterations=args.iterations,
                    device=device,
                )
                all_results.append(r)

        print_results_table(all_results)

        if args.output:
            with open(args.output, 'w') as f:
                json.dump([asdict(r) for r in all_results], f, indent=2, default=str)
            print(f"  Results saved: {args.output}")

    elif args.mode == 'baseline':
        all_results = []
        for bs in batch_sizes:
            for sl in seq_lens:
                print(f"\n  Benchmarking baseline model: bs={bs}, seq={sl}")
                r = benchmark_baseline_model(
                    batch_size=bs, seq_len=sl,
                    hidden_size=args.hidden_size,
                    num_heads=args.num_heads,
                    H_cycles=args.H_cycles, L_cycles=args.L_cycles,
                    H_layers=args.H_layers, L_layers=args.L_layers,
                    warmup=args.warmup, iterations=args.iterations,
                    device=device,
                )
                all_results.append(r)

        print_results_table(all_results)

        if args.output:
            with open(args.output, 'w') as f:
                json.dump([asdict(r) for r in all_results], f, indent=2, default=str)
            print(f"  Results saved: {args.output}")

    print("\n  Done!\n")


if __name__ == '__main__':
    main()