From 2c80ecdf266991c9b231582dfd45fdf82d218cfc Mon Sep 17 00:00:00 2001 From: Parth Ashwin Jain Date: Fri, 24 Jul 2026 03:27:14 -0700 Subject: [PATCH 1/2] [bench] vit_attention: model-matching V stride + warm (no-flush) mode Add --v-stride-mult (default 3) so V is a non-contiguous slice of a packed (B,S,mult,H,D) tensor with token stride mult*H*D, matching a real fused-QKV projection (mult=1 restores the old contiguous layout). Add --no-flush to measure with a warm cache instead of do_bench's cold-cache L2 flush. Emit v_stride_mult / v_token_stride / flush in the JSON config. Co-authored-by: Cursor Signed-off-by: Parth Ashwin Jain --- benchmarks/vit_attention/bench.py | 99 ++++++++++++++++++++++++++++--- 1 file changed, 91 insertions(+), 8 deletions(-) diff --git a/benchmarks/vit_attention/bench.py b/benchmarks/vit_attention/bench.py index f6889efc4a94..5cb367fc89a8 100755 --- a/benchmarks/vit_attention/bench.py +++ b/benchmarks/vit_attention/bench.py @@ -20,6 +20,14 @@ Triton-specific tuning knobs (--bm/--bn/--nw/--ns/--we) are passed directly to the kernel launch and are ignored for other backends. +Operand layout: a real ViT fuses the QKV projection, so V is a *non-contiguous* +view into a packed tensor (token stride = mult*H*D, mult=3 for a standard fused +QKV) while Q/K are made contiguous by the serving wrappers. --v-stride-mult +controls this (default 3 = model-matching fused QKV; 1 = fully contiguous). + +--no-flush measures with a warm cache (no L2 clear between iterations) instead +of do_bench's default cold-cache behavior. + Output: one JSON line per backend to stdout. """ @@ -86,11 +94,66 @@ def _parse() -> argparse.Namespace: p.add_argument( "--we", type=int, default=None, help="waves_per_eu (TRITON_ATTN only)" ) + p.add_argument( + "--v-stride-mult", + type=int, + default=3, + help=( + "V token-stride multiplier: V is a non-contiguous slice of a packed " + "(B,S,mult,H,D) tensor, so its token stride is mult*H*D. Default 3 " + "matches a fused QKV projection (the real model layout); 1 = the " + "old fully-contiguous layout." + ), + ) + p.add_argument( + "--no-flush", + action="store_true", + help=( + "Measure with a warm cache (no L2 flush between iterations) instead " + "of do_bench's default cold-cache flush." + ), + ) p.add_argument("--warmup-ms", type=int, default=200) p.add_argument("--rep-ms", type=int, default=600) return p.parse_args() +def _do_bench_no_flush(fn, warmup_ms: int, rep_ms: int) -> list[float]: + """do_bench variant that does NOT clear the cache between iterations. + + Mirrors triton.testing.do_bench's event-timed loop (5-call cost estimate -> + warmup -> timed reps) but omits the L2 flush, so q/k/v stay resident (warm). + Returns per-call times in milliseconds. + """ + di = triton.runtime.driver.active.get_device_interface() + + fn() + di.synchronize() + + start = di.Event(enable_timing=True) + end = di.Event(enable_timing=True) + start.record() + for _ in range(5): + fn() + end.record() + di.synchronize() + estimate_ms = start.elapsed_time(end) / 5 + n_warmup = max(1, int(warmup_ms / estimate_ms)) + n_repeat = max(1, int(rep_ms / estimate_ms)) + + for _ in range(n_warmup): + fn() + starts = [di.Event(enable_timing=True) for _ in range(n_repeat)] + ends = [di.Event(enable_timing=True) for _ in range(n_repeat)] + di.synchronize() + for i in range(n_repeat): + starts[i].record() + fn() + ends[i].record() + di.synchronize() + return [s.elapsed_time(e) for s, e in zip(starts, ends)] + + def _bench_one( backend_name: str, args: argparse.Namespace, @@ -209,13 +272,17 @@ def _fn(): else: raise ValueError(f"Unsupported backend: {backend_name}") - # do_bench clears the L2 cache before each measurement iteration. - all_times = triton.testing.do_bench( - _fn, - warmup=args.warmup_ms, - rep=args.rep_ms, - return_mode="all", - ) + if args.no_flush: + # Warm cache: no L2 flush between iterations. + all_times = _do_bench_no_flush(_fn, args.warmup_ms, args.rep_ms) + else: + # do_bench clears the L2 cache before each measurement iteration. + all_times = triton.testing.do_bench( + _fn, + warmup=args.warmup_ms, + rep=args.rep_ms, + return_mode="all", + ) all_times = sorted(all_times) n = len(all_times) mean = sum(all_times) / n @@ -231,6 +298,9 @@ def _fn(): "D": D, "dtype": args.dtype, "backend": backend_name, + "v_stride_mult": args.v_stride_mult, + "v_token_stride": int(v4.stride(-3)), + "flush": not args.no_flush, } if backend == AttentionBackendEnum.TRITON_ATTN: config.update( @@ -271,9 +341,22 @@ def main() -> int: # Wrappers (FLASH_ATTN, ROCM_AITER_FA, TRITON_ATTN) expect (B, S, H, D). # TORCH_SDPA expects the same via vit_torch_sdpa_wrapper. + # + # Q/K are contiguous (the serving wrappers make them so). V models the fused + # QKV projection: it is a non-contiguous slice of a packed (B,S,mult,H,D) + # tensor, giving a token stride of mult*H*D (mult=3 => the real fused-QKV + # layout). --v-stride-mult 1 restores the old fully-contiguous V. q4 = torch.randn(B, S, H, D, dtype=dtype, device=device) k4 = torch.randn(B, S, H, D, dtype=dtype, device=device) - v4 = torch.randn(B, S, H, D, dtype=dtype, device=device) + mult = args.v_stride_mult + if mult <= 1: + v4 = torch.randn(B, S, H, D, dtype=dtype, device=device) + else: + # Packed buffer; take index 0 along the fused dim -> V keeps a token + # stride of mult*H*D. The (B,S) dims stay mergeable, so the TRITON path's + # .view(B*S, H, D) still works while preserving the non-contiguous stride. + v_packed = torch.randn(B, S, mult, H, D, dtype=dtype, device=device) + v4 = v_packed[:, :, 0] # cu_seqlens / max_seqlen used by FA and Triton backends cu = torch.tensor([i * S for i in range(B + 1)], dtype=torch.int32, device=device) From a700083800a9b305e075bfcd88b833df35867ea7 Mon Sep 17 00:00:00 2001 From: Parth Ashwin Jain Date: Thu, 30 Jul 2026 00:24:51 -0700 Subject: [PATCH 2/2] [bench] vit_attention: add controllable --flush-mib cache-eviction knob Add --flush-mib N to model a controllable amount of cache eviction by zeroing an N-MiB device buffer before each timed rep (N=0 warm, N~LLC serving-like partial, N>=2x LLC cold), generalizing the binary --no-flush. Also emit flush_mode/flush_mib and v_contiguous in the JSON config. Co-authored-by: Cursor Signed-off-by: Parth Ashwin Jain --- benchmarks/vit_attention/bench.py | 67 ++++++++++++++++++++++++++----- 1 file changed, 56 insertions(+), 11 deletions(-) diff --git a/benchmarks/vit_attention/bench.py b/benchmarks/vit_attention/bench.py index 5cb367fc89a8..f78f70102915 100755 --- a/benchmarks/vit_attention/bench.py +++ b/benchmarks/vit_attention/bench.py @@ -25,8 +25,18 @@ QKV) while Q/K are made contiguous by the serving wrappers. --v-stride-mult controls this (default 3 = model-matching fused QKV; 1 = fully contiguous). ---no-flush measures with a warm cache (no L2 clear between iterations) instead -of do_bench's default cold-cache behavior. +Cache working-set state is controllable, because it dominates the gap vs real +serving (where surrounding qkv-proj / MLP / norm work only *partially* evicts +q/k/v between attention calls): + * default -> stock do_bench (fully cold; clears the L2 each rep) + * --no-flush -> fully warm (q/k/v stay resident between reps) + * --flush-mib N -> zero an N-MiB device buffer before each rep, streaming + it through the last-level cache to displace q/k/v. + N=0 is warm; N ~ LLC size models serving-like partial + eviction; N >= 2x LLC is fully cold. (A write equal to + the cache size does not fully evict due to set/way + associativity + the flush stream evicting itself, so + full-cold needs ~2x the cache size.) Output: one JSON line per backend to stdout. """ @@ -110,7 +120,19 @@ def _parse() -> argparse.Namespace: action="store_true", help=( "Measure with a warm cache (no L2 flush between iterations) instead " - "of do_bench's default cold-cache flush." + "of do_bench's default cold-cache flush. Equivalent to --flush-mib 0." + ), + ) + p.add_argument( + "--flush-mib", + type=int, + default=None, + help=( + "Zero an N-MiB device buffer before each timed rep to model a " + "controllable amount of cache eviction. N=0 = warm (no flush); " + "N ~ last-level-cache size = serving-like partial eviction; " + "N >= 2x LLC = fully cold. Overrides --no-flush. When unset, the " + "default stock do_bench full-cold L2 flush is used." ), ) p.add_argument("--warmup-ms", type=int, default=200) @@ -118,15 +140,24 @@ def _parse() -> argparse.Namespace: return p.parse_args() -def _do_bench_no_flush(fn, warmup_ms: int, rep_ms: int) -> list[float]: - """do_bench variant that does NOT clear the cache between iterations. +def _do_bench_flush(fn, warmup_ms: int, rep_ms: int, flush_mib: int) -> list[float]: + """do_bench variant with a controllable N-MiB cache flush between reps. Mirrors triton.testing.do_bench's event-timed loop (5-call cost estimate -> - warmup -> timed reps) but omits the L2 flush, so q/k/v stay resident (warm). - Returns per-call times in milliseconds. + warmup -> timed reps). Before each timed rep it zeroes a ``flush_mib``-MiB + device buffer; those writes stream through the last-level cache and displace + the resident q/k/v. ``flush_mib=0`` performs no flush (fully warm), so q/k/v + stay resident; ``flush_mib ~ LLC`` models serving-like partial eviction and + ``flush_mib >= 2x LLC`` is fully cold. Returns per-call times in ms. """ di = triton.runtime.driver.active.get_device_interface() + flush_buf = None + if flush_mib > 0: + flush_buf = torch.empty( + int(flush_mib) * 1024 * 1024, dtype=torch.int8, device="cuda" + ) + fn() di.synchronize() @@ -147,6 +178,8 @@ def _do_bench_no_flush(fn, warmup_ms: int, rep_ms: int) -> list[float]: ends = [di.Event(enable_timing=True) for _ in range(n_repeat)] di.synchronize() for i in range(n_repeat): + if flush_buf is not None: + flush_buf.zero_() starts[i].record() fn() ends[i].record() @@ -272,9 +305,16 @@ def _fn(): else: raise ValueError(f"Unsupported backend: {backend_name}") - if args.no_flush: - # Warm cache: no L2 flush between iterations. - all_times = _do_bench_no_flush(_fn, args.warmup_ms, args.rep_ms) + if args.flush_mib is not None: + # Controllable partial eviction: zero an N-MiB buffer before each rep. + all_times = _do_bench_flush(_fn, args.warmup_ms, args.rep_ms, args.flush_mib) + flush_mode = "warm" if args.flush_mib == 0 else "flush_mib" + flush_mib = args.flush_mib + elif args.no_flush: + # Warm cache: no L2 flush between iterations (== --flush-mib 0). + all_times = _do_bench_flush(_fn, args.warmup_ms, args.rep_ms, 0) + flush_mode = "warm" + flush_mib = 0 else: # do_bench clears the L2 cache before each measurement iteration. all_times = triton.testing.do_bench( @@ -283,6 +323,8 @@ def _fn(): rep=args.rep_ms, return_mode="all", ) + flush_mode = "cold" + flush_mib = None all_times = sorted(all_times) n = len(all_times) mean = sum(all_times) / n @@ -300,7 +342,10 @@ def _fn(): "backend": backend_name, "v_stride_mult": args.v_stride_mult, "v_token_stride": int(v4.stride(-3)), - "flush": not args.no_flush, + "v_contiguous": bool(v4.is_contiguous()), + "flush": flush_mode != "warm", + "flush_mode": flush_mode, + "flush_mib": flush_mib, } if backend == AttentionBackendEnum.TRITON_ATTN: config.update(