Skip to content

Fix benchmark exploit via object-identity caching #102

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
msaroufim wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fix benchmark exploit via object-identity caching #102

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
583bf08
Fix benchmark exploit via object-identity caching
msaroufim Feb 7, 2026
81d6d4e
Shuffle iteration data and fix recheck bug
msaroufim Feb 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 20 additions & 8 deletions problems/nvidia/eval_better_bench_grouped_gemm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import base64
import dataclasses
import multiprocessing
import random
import re
import time
import os
Expand Down Expand Up @@ -242,10 +243,14 @@ def _run_single_benchmark(
data_list = []
# generate input data once

local_seed = test.args.get("seed", None)
for i in range(NUM_ITERATIONS_PER_BENCHMARK):
if "seed" in test.args:
test.args["seed"] += 42
data = generate_input(**test.args)
if local_seed is not None:
local_seed += 42
args = {**test.args, "seed": local_seed}
else:
args = test.args
data = generate_input(**args)
data_list.append(data)

check_copy = _clone_data(data_list)
Expand All @@ -270,14 +275,21 @@ def _run_single_benchmark(

bm_start_time = time.perf_counter_ns()
for i in range(max_repeats):
# Clone and shuffle data before timing to prevent both
# object-identity caching and call-order caching exploits
iteration_data = _clone_data(data_list)
shuffle_order = list(range(len(iteration_data)))
random.shuffle(shuffle_order)
iteration_data = [iteration_data[j] for j in shuffle_order]

torch.cuda.synchronize()

outputs = []
clear_l2_cache()
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()
for data in data_list:
for data in iteration_data:
output = custom_kernel(data)
outputs.append(output)
end_event.record()
Expand All @@ -287,10 +299,10 @@ def _run_single_benchmark(
) * 1e6 # Convert ms to ns

if recheck:
for reference_output, custom_output in zip(check_copy, outputs):
good, message = check_implementation(reference_output, custom_output)
if not good:
return message
for j, custom_output in zip(shuffle_order, outputs):
good, message = check_implementation(check_copy[j], custom_output)
if not good:
return message

durations.append(duration)

Expand Down