[FA4/B200] return_lse=True makes D64 forward about 19% faster

GPU: NVIDIA B200
Torch: 2.14.0.dev20260722+cu132 (CUDA 13.2)
CuTe DSL: 4.6.0.dev0
FA4 commit: 00756db9d921da0846453283ddfbeb7457abd09b
no-LSE: median=3.0018 ms rounds=[3.007457275390625, 3.0017739868164064, 3.001456604003906, 3.0028912353515627, 3.0008636474609376]
with-LSE: median=2.5281 ms rounds=[2.5091484069824217, 2.545796203613281, 2.512978210449219, 2.5281126403808596, 2.532301177978516]
with-LSE speedup: 1.187x

python repro_fa4_lse_b200.py /path/to/flash-attention

"""Reproduce the FA4 D64 B200 performance inversion caused by returning LSE."""

import argparse
import importlib
import importlib.metadata
import statistics
import subprocess
import sys
import time
from pathlib import Path
from types import ModuleType
from typing import Callable

import torch
import torch.nn.functional as F


def load_fa4_interface(fa4_root: Path) -> ModuleType:
    """Load the CuTe FA4 interface without requiring the FA2 extension."""
    package = ModuleType("flash_attn")
    package.__path__ = [str(fa4_root / "flash_attn")]
    sys.modules["flash_attn"] = package
    return importlib.import_module("flash_attn.cute.interface")


def make_call(
    interface: ModuleType,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    return_lse: bool,
) -> Callable[[], torch.Tensor]:
    """Create one allocation-free FA4 forward specialization."""
    out = torch.empty_like(q)
    lse = (
        torch.empty(q.shape[0], q.shape[2], q.shape[1], device=q.device)
        if return_lse
        else None
    )

    def run() -> torch.Tensor:
        return interface._flash_attn_fwd(
            q, k, v, out=out, lse=lse, return_lse=return_lse
        )[0]

    return run


def capture(function: Callable[[], torch.Tensor]) -> torch.cuda.CUDAGraph:
    """Capture a warmed FA4 call into a fixed-pointer CUDA graph."""
    stream = torch.cuda.Stream()
    stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(stream):
        for _ in range(3):
            function()
    stream.synchronize()
    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph, stream=stream):
        function()
    return graph


def warm_gpu(seconds: float = 1.0) -> None:
    """Warm the GPU long enough to avoid measuring initial idle clocks."""
    x = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
    deadline = time.perf_counter() + seconds
    while time.perf_counter() < deadline:
        torch.mm(x, x)
    torch.cuda.synchronize()


def time_graph(graph: torch.cuda.CUDAGraph, iterations: int) -> float:
    """Return mean fixed-pointer graph replay latency in milliseconds."""
    for _ in range(25):
        graph.replay()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    for _ in range(iterations):
        graph.replay()
    end.record()
    end.synchronize()
    return start.elapsed_time(end) / iterations


def main() -> None:
    """Compile, validate, and compare no-LSE and with-LSE FA4 kernels."""
    parser = argparse.ArgumentParser()
    parser.add_argument("fa4_root", type=Path)
    parser.add_argument("--iterations", type=int, default=100)
    parser.add_argument("--rounds", type=int, default=5)
    args = parser.parse_args()

    torch.manual_seed(0)
    torch.cuda.manual_seed_all(0)
    interface = load_fa4_interface(args.fa4_root.resolve())
    shape = (2, 16384, 16, 64)
    q, k, v = (
        torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)
    )
    no_lse = make_call(interface, q, k, v, False)
    with_lse = make_call(interface, q, k, v, True)

    expected = F.scaled_dot_product_attention(
        q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
    ).transpose(1, 2)
    torch.testing.assert_close(no_lse().float(), expected.float(), atol=5e-2, rtol=2e-2)
    torch.testing.assert_close(
        with_lse().float(), expected.float(), atol=5e-2, rtol=2e-2
    )
    torch.cuda.synchronize()

    graphs = {"no-LSE": capture(no_lse), "with-LSE": capture(with_lse)}
    samples = {name: [] for name in graphs}
    warm_gpu()
    for round_index in range(args.rounds):
        order = tuple(graphs) if round_index % 2 == 0 else tuple(reversed(graphs))
        for name in order:
            samples[name].append(time_graph(graphs[name], args.iterations))

    commit = subprocess.check_output(
        ["git", "-C", str(args.fa4_root), "rev-parse", "HEAD"], text=True
    ).strip()
    medians = {name: statistics.median(values) for name, values in samples.items()}
    print(f"GPU: {torch.cuda.get_device_name()}")
    print(f"Torch: {torch.__version__} (CUDA {torch.version.cuda})")
    print(f"CuTe DSL: {importlib.metadata.version('nvidia-cutlass-dsl')}")
    print(f"FA4 commit: {commit}")
    for name, values in samples.items():
        print(f"{name}: median={medians[name]:.4f} ms rounds={values}")
    print(f"with-LSE speedup: {medians['no-LSE'] / medians['with-LSE']:.3f}x")


if __name__ == "__main__":
    main()