Skip to content

feat(linked): add FlashInfer sampling provider - #930

Draft
voltjia wants to merge 1 commit into
masterfrom
feat/linked-flashinfer-sampling
Draft

feat(linked): add FlashInfer sampling provider#930
voltjia wants to merge 1 commit into
masterfrom
feat/linked-flashinfer-sampling

Conversation

@voltjia

@voltjia voltjia commented Aug 13, 2026

Copy link
Copy Markdown
Collaborator

Summary

Design

The provider composes FlashInfer's TVM FFI entry points:

  • top_k_first: top-k mask -> softmax -> top-p sampling
  • joint: softmax -> joint top-k/top-p sampling

Indices 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_p remains 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.

  • Resolver/schema tests: 53 passed.
  • Release CUDA build: infiniops and Python ops targets passed for SM80.
  • FlashInfer implementation 16 GPU tests: 17 passed.
  • Dynamic dependencies: readelf found sampling.so and libtvm_ffi.so; GPU-mounted ldd -r had no unresolved dependencies.
  • Descriptor cache stress (batch 2..32, vocab 32768): 0 MiB retained device memory after synchronization and allocator cleanup.
  • Ruff 0.15.22 check and format check passed; git diff --check passed.

Tested with flashinfer-jit-cache==0.6.7.post3+cu130, apache-tvm-ffi==0.1.10, CUDA 13.1, and an SM80 target.

Dependencies

@voltjia
voltjia force-pushed the feat/operator-implementation-selection branch from 4ba65ce to 87a242b Compare August 13, 2026 16:43
@voltjia
voltjia force-pushed the feat/linked-flashinfer-sampling branch from 8054804 to 94965ca Compare August 13, 2026 16:43
@voltjia
voltjia force-pushed the feat/operator-implementation-selection branch from 87a242b to 5c7ef80 Compare August 18, 2026 09:34
Base automatically changed from feat/operator-implementation-selection to master August 19, 2026 01:53
@voltjia
voltjia force-pushed the feat/linked-flashinfer-sampling branch from 94965ca to 2e69e64 Compare August 24, 2026 09:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant