[bench] vit_attention/bench.py: add --backend, default to running all

[mgehre-amd]

Add a --backend flag whose choices match the values accepted by --mm-encoder-attn-backend (TRITON_ATTN, TORCH_SDPA, FLASH_ATTN, ROCM_AITER_FA). The default value "all" runs every backend in sequence and emits one JSON line per backend so backends can be compared in a single invocation.

Also sync the TRITON_ATTN launch with the current _fwd_kernel signature, which gained BLOCK_DMODEL_TAIL and HEAD_STRIDE_ALIGNED_8 since the script was first written. Use _split_head_dim() and get_block_n() from the production code so the tuning parameters are consistent.

Run via
benchmarks/vit_attention/bench.py --heads 16 --head-dim 64 --seq 784
to get output like

{"per_call_ms_mean": 0.15117897940499167, "per_call_ms_median": 0.1509760022163391, "per_call_ms_min": 0.14600500464439392, "per_call_ms_p10": 0.1490119993686676, "per_call_ms_p90": 0.15294000506401062, "samples": 248, "config": {"B": 1, "S": 784, "H": 16, "D": 64, "dtype": "bf16", "backend": "TRITON_ATTN", "BLOCK_M": 128, "BLOCK_DMODEL": 64, "BLOCK_DMODEL_TAIL": 0, "BLOCK_N": 128, "num_warps": 8, "num_stages": 1, "waves_per_eu": null}}
{"per_call_ms_mean": 0.23133230818324982, "per_call_ms_median": 0.23059099912643433, "per_call_ms_min": 0.2106119990348816, "per_call_ms_p10": 0.2198610007762909, "per_call_ms_p90": 0.23941799998283386, "samples": 438, "config": {"B": 1, "S": 784, "H": 16, "D": 64, "dtype": "bf16", "backend": "TORCH_SDPA"}}
{"per_call_ms_mean": 0.28636176208245384, "per_call_ms_median": 0.2856689989566803, "per_call_ms_min": 0.244828999042511, "per_call_ms_p10": 0.25616100430488586, "per_call_ms_p90": 0.30823400616645813, "samples": 388, "config": {"B": 1, "S": 784, "H": 16, "D": 64, "dtype": "bf16", "backend": "FLASH_ATTN"}}