[bench] vit_attention: model-matching V stride + controllable cache-eviction (--flush-mib)

[parthash0804]

Summary

Extends the ViT attention microbenchmark (benchmarks/vit_attention/bench.py) so it can faithfully reproduce full-model serving conditions for the ViT-encoder attention path. Two knobs are added:

• --v-stride-mult N (default 3) — model the real operand memory layout.
  In serving, V is a non-contiguous slice of the fused QKV projection
  (token stride = N*H*D, i.e. 3072 for the real model), while Q/K are made
  contiguous by the serving wrappers. --v-stride-mult 3 reproduces this;
  1 restores the old fully-contiguous V.
• --flush-mib N — model a controllable amount of cache eviction between
  timed reps, generalizing the previous binary warm/cold behavior:
  • default (unset) → stock do_bench full-cold L2 flush
  • --no-flush (== --flush-mib 0) → fully warm (q/k/v stay resident)
  • --flush-mib N → zero an N-MiB device buffer before each rep; N ~ LLC
    models serving-like partial eviction, N >= 2x LLC is fully cold

Why

The single-op microbench and the in-serving profiler were not measuring the same thing, which made per-call ViT attention latencies look inconsistent. Root-cause showed the gap is fully explained by measurable effects — dominant one being cache working-set state (do_bench is fully cold; serving is only partially evicted between attention calls), plus the non-contiguous
fused-QKV V operand layout. These flags let the microbench match serving:

• On Strix Halo (gfx1151, 32 MB MALL): a single-pass --flush-mib 32 (≈ LLC size) reproduces serving to ~1% across fp16 and AWQ.
• On Strix Point (gfx1150, 2 MB L2): with the clock pinned, the fused-QKV stride (--v-stride-mult 3) + cold flush lands within a few % of serving.

Changes

• --v-stride-mult builds V as a non-contiguous slice of a packed (B,S,mult,H,D) tensor (token stride mult*H*D); mult=1 = contiguous.
• New _do_bench_flush(fn, warmup, rep, flush_mib) — do_bench-style event-timed loop that zeroes an N-MiB buffer before each rep (flush_mib=0 = warm).
• --no-flush retained as an alias for --flush-mib 0.
• JSON config now reports v_stride_mult, v_token_stride, v_contiguous, flush, flush_mode (cold/warm/flush_mib), and flush_mib.

Results

• python benchmarks/vit_attention/bench.py --backend TRITON_ATTN --seq 1408 --heads 16 --head-dim 64 --dtype fp16 --v-stride-mult 3 --no-flush --bm 128 --bn 128 --nw 8

{"per_call_ms_mean": 0.37667454618736146, "per_call_ms_median": 0.37690699100494385, "per_call_ms_min": 0.3467699885368347, "per_call_ms_p10": 0.3615579903125763, "per_call_ms_p90": 0.3880069851875305, "samples": 1477, "config": {"B": 1, "S": 1408, "H": 16, "D": 64, "dtype": "fp16", "backend": "TRITON_ATTN", "v_stride_mult": 3, "v_token_stride": 3072, "v_contiguous": false, "flush": false, "flush_mode": "warm", "flush_mib": 0, "BLOCK_M": 128, "BLOCK_DMODEL": 64, "BLOCK_DMODEL_TAIL": 0, "BLOCK_N": 128, "num_warps": 8, "num_stages": 1, "waves_per_eu": null}}

[mgehre-amd]

Thanks!