flashinfer / wrapperb91d34
flashinfer_wrapper_b91d34 · flashinfer · 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-b91d34?include=source"interfacepython
revisionda915083d4c7
symbolrun
pathmain.py
Compatibility
declared hardwareNVIDIA A100, NVIDIA B200, 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:d0336f7223c83cf443c5336a5abe3f37e6b82649f88e24f20db1b65aad681a0b
license declaredApache-2.0
license concludedApache-2.0
authorsflashinfer
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 == 202048
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