feat(linked): add FlashInfer sampling provider - #930
Draft
voltjia wants to merge 1 commit into
Draft
Conversation
voltjia
force-pushed
the
feat/operator-implementation-selection
branch
from
August 13, 2026 16:43
4ba65ce to
87a242b
Compare
voltjia
force-pushed
the
feat/linked-flashinfer-sampling
branch
from
August 13, 2026 16:43
8054804 to
94965ca
Compare
This was referenced Aug 13, 2026
voltjia
force-pushed
the
feat/operator-implementation-selection
branch
from
August 18, 2026 09:34
87a242b to
5c7ef80
Compare
Base automatically changed from
feat/operator-implementation-selection
to
master
August 19, 2026 01:53
19 tasks
voltjia
force-pushed
the
feat/linked-flashinfer-sampling
branch
from
August 24, 2026 09:50
94965ca to
2e69e64
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
TopKTopPSamplingFromLogits, backed by FlashInfer's AOT sampling library through the TVM FFI transport merged in feat(linked): support TVM FFI providers #945.fp16,bf16, andfp32logits;top_k_firstandjoint; CPU/CUDA int32/int64 indices; deterministic seed/offset; caller-provided workspace; and non-default CUDA streams.flashinfer-jit-cache>=0.6.7,<0.7andapache-tvm-ffi==0.1.10.ops.json#931 and feat(linked): support TVM FFI providers #945.Design
The provider composes FlashInfer's TVM FFI entry points:
top_k_first: top-k mask -> softmax -> top-p samplingjoint: softmax -> joint top-k/top-p samplingIndices are gathered into a contiguous FP32 logits matrix before sampling. Float64 logits are rejected explicitly because this provider cannot preserve their ordering through that FP32 matrix. Float64
top_premains accepted and is converted toward the interior of(0, 1)so underflow or rounding cannot accidentally disable top-p filtering.Two pinned host staging slots and CUDA events preserve asynchronous H2D staging. A missing caller workspace is allocated and released per call with
cudaMallocAsync/cudaFreeAsync; cached descriptors therefore do not retain shape-sized GPU matrices. CUDA, TVM FFI, metadata, and unsupported-input checks remain active in Release builds.Validation
Manual NVIDIA validation used the exact #945 tree plus the six blobs from commit
2e69e648; post-#945 master changes are unrelated backend work and are covered by PR CI.infiniopsand Pythonopstargets passed for SM80.readelffoundsampling.soandlibtvm_ffi.so; GPU-mountedldd -rhad no unresolved dependencies.git diff --checkpassed.Tested with
flashinfer-jit-cache==0.6.7.post3+cu130,apache-tvm-ffi==0.1.10, CUDA 13.1, and an SM80 target.Dependencies
ops.json#931 (merged): selects exact implementation slots through consumer-ownedops.json.