onnx export with dynamic shapes, fast attention

[jpata]

• factorize ONNX exporting to notebooks/cms/cms-validate-onnx.py which contains a minimal, standalone version of the attention-based model, validated against the base model, and the exporting code to ONNX, with both unfused (default, slow) and fused (fast on GPU) attention
• dynamic shapes: ONNX export now works through TorchScript (the old, non-beta way in pytorch)
• fused attention: we override the aten::scaled_dot_product_attention op to use com.microsoft.MultiHeadAttention from ONNX contrib
• it's fast: on A100, the MultiHeadAttention op executes efficient flash attention kernels, such that the whole model runs in about 15ms and <2GB for 5000 inputs: Integrate new pytorch attention model in CMSSW via ONNX #216 (comment)
• the outputs are equivalent in terms of physics to the original model
• added to huggingface in https://huggingface.co/jpata/particleflow/commit/65729d798b0e51598916b4ae9c8c4f712820c79d

Here's how the direct export of torch.nn.functional.scaled_dot_product_attention to an unfused ONNX model, with full matrix multiplications looks like:

Using the SDPA fused operation that will use flash attention on sufficiently new GPUs, where the MatMul->Softmax->MatMul part in the very end is rolled into an op SDPA that calls MultiHeadAttention:

Here are the timings, showing the benefit of the fused model:

timing/gpu_fp32_fused.txt:Nelem=2560 mean_time=6.99 ms stddev_time=2.89 ms mem_used=1678 MB
timing/gpu_fp32_fused.txt:Nelem=5120 mean_time=16.59 ms stddev_time=0.15 ms mem_used=1946 MB
timing/gpu_fp32_fused.txt:Nelem=10240 mean_time=53.13 ms stddev_time=0.23 ms mem_used=1946 MB

timing/gpu_fp32_unfused.txt:Nelem=2560 mean_time=39.31 ms stddev_time=1.73 ms mem_used=3817 MB
timing/gpu_fp32_unfused.txt:Nelem=5120 mean_time=130.18 ms stddev_time=6.52 ms mem_used=12407 MB
timing/gpu_fp32_unfused.txt:Nelem=10240 mean_time=465.09 ms stddev_time=25.82 ms mem_used=46766 MB