flashinfer_wrapper_4ec4ec35
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-4ec4ec35?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, int32
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:d137fddf6793bc682ba5e1935e706bf3e02edf4ac28b2011fcbf3e0691fc49c4
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_k):
batch_size, vocab_size = probs.shape
device = probs.device
assert vocab_size == 129280
probs = probs.to(torch.float32)
samples = flashinfer.sampling.top_k_sampling_from_probs(
probs=probs,
top_k=top_k,
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