MCPcopy Index your code
hub / github.com/FlashSampling/FlashSampling / nvbench_kernel

Function nvbench_kernel

src/fused_mm_sampling/bench/speed_test.py:161–177  ·  view source on GitHub ↗
(state: "nvbench.State")

Source from the content-addressed store, hash-verified

159 """Run benchmarks using NVBench."""
160
161 def nvbench_kernel(state: "nvbench.State"):
162 provider = state.get_string("Provider")
163 case = as_case(args, provider)
164 kwargs = case.make_fn_kwargs()
165 sampler = get_sampler(provider, weights=kwargs["weights"])
166 sampler.prepare()
167
168 # Warmup (compile, autotune, etc.)
169 sampler.sample(**kwargs)
170 torch.cuda.synchronize()
171
172 def launcher(launch: "nvbench.Launch"):
173 stream = _as_torch_stream(launch.get_stream())
174 with torch.cuda.stream(stream):
175 sampler.sample(**kwargs)
176
177 state.exec(launcher, batched=False)
178
179 csv_args = []
180 if args.tgt_dir is not None:

Callers

nothing calls this directly

Calls 5

as_caseFunction · 0.85
get_samplerFunction · 0.85
make_fn_kwargsMethod · 0.80
prepareMethod · 0.45
sampleMethod · 0.45

Tested by

no test coverage detected