flashinfer_wrapper_d86b24bd
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-d86b24bd?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:deaec74210fbe8b3fe27d70bf2769233df186e79388ea882456cbf95c2755a5a
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 == 128256
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