Function _run_sampling_native_workload

benchmarks/profiler.py:684–718

(
    profiler: PerformanceProfiler,
    args: argparse.Namespace,
    device: torch.device,
    dtype: torch.dtype,
    batch_size: int,
    seq_len: int,
    vocab_size: int,
)

def _run_sampling_native_workload(
    profiler: PerformanceProfiler,
    args: argparse.Namespace,
    device: torch.device,
    dtype: torch.dtype,
    batch_size: int,
    seq_len: int,
    vocab_size: int,
) -> BenchmarkMetrics:
    try:
        fn = _sampling_fn(
            batch_size=batch_size,
            vocab_size=vocab_size,
            dtype=torch.float32 if device.type == "cpu" else dtype,
            device=device,
            seed=args.seed + 2,
            top_k=args.top_k,
            top_p=args.top_p,
        )
        return profiler.profile_sampling(
            fn,
            batch_size=batch_size,
            seq_len=seq_len,
            vocab_size=vocab_size,
            benchmark_name="sampling_native",
        )
    except Exception as exc:
        return _blocked_metric(
            profiler=profiler,
            benchmark_name="sampling_native",
            batch_size=batch_size,
            seq_len=seq_len,
            vocab_size=vocab_size,
            exc=exc,
        )


WORKLOAD_REGISTRY: dict[str, WorkloadRunner] = {