flashinfer_wrapper_5df4fa0b
FlashInfer-Bench baselines · python · Apache-2.0
Use it
Vendorable · source mirrored · Apache-2.0View source →
No package. Vendor the mirrored source: 25 lines, Apache-2.0, pinned at da91508.
main.py
curl "https://kernelindex.com/api/v1/implementations/flashinfer-flashinfer-wrapper-5df4fa0b?include=source"interfacepython
revisionda915083d4c7
symbolrun
pathmain.py
Compatibility
declared hardwareNVIDIA A100, NVIDIA B200, NVIDIA GeForce RTX 4090, NVIDIA H100, NVIDIA H20, NVIDIA H200
architecturesunknown
dtypesfp32
Benchmark evidence
No published measurement for this revision.
No evidence · How evidence levels are derived →
Source and license
sourcehttps://huggingface.co/datasets/flashinfer-ai/flashinfer-trace
commitda915083d4c7c5e61aa3005e3d17ae488e0fc71c
revision digestsha256:f4f7ba47ca1d2deac0096d64a801f2f6b3bdde053960b053ac2680834184bd1c
license declaredApache-2.0
license concludedApache-2.0
authorsbaseline
imported2026-08-16
Kernel source
main.py25 lines
import torch
import flashinfer
def run(probs, top_p):
batch_size, vocab_size = probs.shape
device = probs.device
assert vocab_size == 128256
probs = probs.to(torch.float32)
samples = flashinfer.sampling.top_p_sampling_from_probs(
probs=probs,
top_p=top_p,
indices=None,
deterministic=False,
generator=None,
check_nan=False
)
samples = samples.to(torch.int64)
return samples
Source code from the importing source · Apache-2.0
No published measurement for this revision
JSON