File size: 780 Bytes
b22e03e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | # 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:
```bash
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
```
|