flashinfer_wrapper_9c1e50fa
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-9c1e50fa?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:ff4e25a3a86ce9b72de99ff82440537bba6be82fa93dd44aad436d7ab1b79a2a
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 == 151936
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