FlashAttentionの速度とGPUメモリを比較する方法｜PyTorchの測定コードと注意点

"""Compare two SDPA forward backends on the same CUDA tensors.

Run: python benchmark.py > results.jsonl
This example has not been executed on a Flash-capable GPU by the author.
"""
from __future__ import annotations

import json
import logging
import statistics
import sys

try:
    import torch
    import torch.nn.functional as F
    from torch.nn.attention import SDPBackend, sdpa_kernel
except ImportError:
    print('ERROR: Install a CUDA-enabled PyTorch with torch.nn.attention.sdpa_kernel.', file=sys.stderr)
    raise SystemExit(2)

logging.basicConfig(level=logging.INFO, format='%(levelname)s %(message)s')
LOGGER = logging.getLogger(__name__)
BACKENDS = {'math': SDPBackend.MATH, 'flash': SDPBackend.FLASH_ATTENTION}


def emit(row: dict[str, object]) -> None:
    """Write one record. Args: row, JSON-compatible fields. Returns: None.

    Raises: TypeError for non-JSON values. Example: emit({'status': 'ok'}).
    """
    print(json.dumps(row, ensure_ascii=False), flush=True)


def forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    """Evaluate square causal attention. Args: q/k/v, CUDA B,H,N,D tensors.

    Returns: output tensor. Raises: RuntimeError for unsupported inputs.
    Example: output = forward(q, k, v).
    """
    # Explicit dropout avoids changing the operation between comparisons.
    return F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=True)


def measure(name: str, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> dict[str, object]:
    """Measure a warmed forward. Args: backend name and identical input tensors.

    Returns: latency and allocator statistics. Raises: RuntimeError/OOM.
    Example: row = measure('math', q, k, v).
    """
    with sdpa_kernel([BACKENDS[name]]):
        for _ in range(10):
            out = forward(q, k, v)
            del out  # Do not keep an earlier result alive during the next call.
        torch.cuda.synchronize()
        samples: list[float] = []
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        for _ in range(5):
            start.record()
            for _ in range(20):
                out = forward(q, k, v)
                del out
            end.record()
            end.synchronize()
            samples.append(start.elapsed_time(end) / 20)

        # Measure memory separately so reference outputs cannot inflate the baseline.
        torch.cuda.synchronize()
        baseline = torch.cuda.memory_allocated()
        torch.cuda.reset_peak_memory_stats()
        out = forward(q, k, v)
        torch.cuda.synchronize()
        peak = torch.cuda.max_memory_allocated()
        del out

    return {
        'backend': name, 'status': 'ok',
        'median_ms': statistics.median(samples),
        'min_ms': min(samples), 'max_ms': max(samples),
        'baseline_MiB': baseline / 2**20,
        'peak_allocated_MiB': peak / 2**20,
        'extra_peak_MiB': (peak - baseline) / 2**20,
    }


@torch.inference_mode()
def run_case(n: int) -> None:
    """Compare one sequence length. Args: n, positive sequence length.

    Returns: None; writes JSON lines. Raises: RuntimeError on input allocation.
    Example: run_case(512).
    """
    shape = (1, 8, n, 64)
    q, k, v = [torch.randn(shape, device='cuda', dtype=torch.float16) for _ in range(3)]
    ok: set[str] = set()
    for name in BACKENDS:
        LOGGER.info('N=%d backend=%s', n, name)
        try:
            row = measure(name, q, k, v)
            ok.add(name)
        except torch.cuda.OutOfMemoryError as error:
            row = {'backend': name, 'status': 'oom', 'reason': str(error)}
        except RuntimeError as error:
            # Report the real error: not every RuntimeError means unsupported hardware.
            row = {'backend': name, 'status': 'error', 'reason': str(error)}
        emit({'kind': 'measurement', 'N': n, **row})
        torch.cuda.empty_cache()  # Only between backends, never inside the timed section.

    if ok == set(BACKENDS):
        try:
            # Accuracy checks follow both measurements to avoid retaining a reference tensor.
            with sdpa_kernel([SDPBackend.MATH]):
                reference = forward(q, k, v)
            with sdpa_kernel([SDPBackend.FLASH_ATTENTION]):
                actual = forward(q, k, v)
            diff = (actual.float() - reference.float()).abs()
            finite = bool(torch.isfinite(reference).all() & torch.isfinite(actual).all())
            emit({'kind': 'output_check', 'N': n,
                  'finite': finite,
                  'max_abs_diff': diff.max().item() if finite else None,
                  'mean_abs_diff': diff.mean().item() if finite else None})
        except RuntimeError as error:
            emit({'kind': 'output_check', 'N': n, 'status': 'error', 'reason': str(error)})


def main() -> int:
    """Run the example. Args: none. Returns: 0, or 2 without CUDA.

    Raises: unexpected non-RuntimeError failures. Example: raise SystemExit(main()).
    """
    if not torch.cuda.is_available():
        LOGGER.error('CUDA is unavailable; no performance measurements were made.')
        return 2
    torch.manual_seed(1234)
    emit({'kind': 'environment', 'torch': torch.__version__,
          'cuda_runtime': torch.version.cuda, 'gpu': torch.cuda.get_device_name(0),
          'capability': torch.cuda.get_device_capability(0),
          'B': 1, 'H': 8, 'D': 64, 'dtype': 'float16', 'causal': True,
          'dropout': 0.0, 'mode': 'inference_forward', 'warmup': 10,
          'repeats_per_interval': 20, 'intervals': 5})
    for n in (512, 1024, 2048):
        try:
            run_case(n)
        except RuntimeError as error:
            emit({'kind': 'case', 'N': n, 'status': 'error', 'reason': str(error)})
        torch.cuda.empty_cache()
    return 0


if __name__ == '__main__':
    raise SystemExit(main())

python benchmark.py > results.jsonl 2> benchmark.log