Skip to content

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

Draft
voltjia wants to merge 2 commits into
feat/operator-implementation-selectionfrom
feat/linked-flashinfer-sampling
Draft

feat(linked): add FlashInfer sampling provider#930
voltjia wants to merge 2 commits into
feat/operator-implementation-selectionfrom
feat/linked-flashinfer-sampling

Conversation

@voltjia

@voltjia voltjia commented Aug 13, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Adds a generic TVM-FFI linked-provider transport with installed header, dependent-library, CUDA source, RPATH, and Python distribution version resolution.
  • Adds NVIDIA implementation slot 16 for TopKTopPSamplingFromLogits backed by FlashInfer's AOT sampling module.
  • Supports fp16, bf16, fp32, and fp64 logits; top_k_first and joint; CPU/CUDA int32/int64 indices; deterministic seed/offset; default or caller workspace; and non-default CUDA streams.
  • Requires flashinfer-jit-cache>=0.6.7,<0.7 and apache-tvm-ffi==0.1.10.

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, preserving repeated/reordered indices and output batches that differ from the unique-logits batch. Two pinned host staging slots and CUDA events preserve asynchronous stream semantics; the provider-owned default workspace is sequenced across streams, while caller-provided workspace follows the existing Handle ownership contract.

The linked DSOs remain external dependencies under the existing linked-provider deployment contract. They are resolved at configure time and referenced through runtime paths; this PR does not bundle FlashInfer or TVM-FFI into InfiniOps wheels.

Validation

  • Release/NDEBUG NVIDIA build passed without provider warnings on A100.
  • Linked resolver tests: 33 passed.
  • FlashInfer slot 16 focused GPU tests: 10 passed.
  • Coverage includes both filter orders, host/CUDA indices, int32/int64 output, seed/offset, default workspace, and a non-default-stream nonblocking sentinel.
  • Ruff format/check, clang-format 21.1.8 dry-run, and git diff --check passed.
  • Hosted build, Ruff, and clang-format checks passed.
  • FlashInfer packed ABI was checked from 0.6.7 through 0.6.18rc1; 0.6.6 is incompatible and rejected by the configured PEP 440 constraint.

Known Boundary

An isolated C++ exact-external-workspace probe compiled, but its test container failed at the first standalone cudaMalloc before entering InfiniOps. The Python GPU matrix exercised the provider-owned workspace and stream path; exact external-workspace runtime remains a residual validation gap rather than an observed provider failure.

Dependency

Stacked on #931, which adds exact implementation selection through consumer-owned ops.json. After #931 merges, this PR can be rebased and retargeted to master.

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