failing example

gdn_prefill_qk4_v8_d128_k_last.py

from flashinfer_bench import TraceSet, BenchmarkConfig
from flashinfer_bench.bench.utils import gen_inputs, load_safetensors
from flashinfer_bench.bench.evaluators.default import DefaultEvaluator
from flashinfer_bench.compile import BuilderRegistry
import torch

ts = TraceSet.from_path("./data/mlsys26-contest")
defn = ts.definitions["gdn_prefill_qk4_v8_d128_k_last"]
workload = [
    w
    for w in ts.workloads["gdn_prefill_qk4_v8_d128_k_last"]
    if w.workload.axes["total_seq_len"] == 83
][0].workload

ref = BuilderRegistry.get_instance().build_reference(defn)
safe = load_safetensors(defn, workload, ts.root)
inputs = gen_inputs(defn, workload, device="cuda", safe_tensors=safe)

with torch.no_grad():
    output, new_state = ref(*inputs)

print(f"reference output has inf: {torch.isinf(output).any().item()}")
print(f"reference output has nan: {torch.isnan(output).any().item()}")

# Run check_correctness with the reference as BOTH the solution and the reference
cfg = BenchmarkConfig()
correctness, eval_result = DefaultEvaluator.check_correctness(
    definition=defn,
    sol_runnable=ref,
    inputs=[inputs],
    ref_outputs=[[output, new_state]],
    cfg=cfg,
    log_path="",
    device="cuda",
)

print(
    f"\ncheck_correctness result: {eval_result.status.value if eval_result else 'PASSED'}"
)
print(f"  abs_err={correctness.max_absolute_error}")
print(f"  rel_err={correctness.max_relative_error}")

reference output has inf: True
reference output has nan: True

# check_correctness result: INCORRECT_NUMERICAL
#   abs_err=inf
#   rel_err=inf