Benchmarks

benchmark.py compares the package API against the current two-call SDPA baseline at PI052 FlexAttention replacement shapes. It records forward time, forward+backward time, peak memory, forward max absolute diff, and gradient norm relative diff.

The default shape is the measured PI052 flow-only shape:

B=4
heads=8
head_dim=256
prefix_len=700
action_blocks=5
action_block_size=50

For hardware/tile sweeps, run one process per shape and preset to avoid torch.compile guard reuse noise:

python flashrt-flex-attention-train/benchmarks/shape_matrix.py --presets a100 --output /tmp/a100_flex_matrix.jsonl
python flashrt-flex-attention-train/benchmarks/shape_matrix.py --presets consumer --output /tmp/consumer_flex_matrix.jsonl