FlashInfer sampling - #2456

Merged

tdene merged 23 commits into

NVIDIA:mainfrom

tdene:tde/flashinfer_sampling

May 5, 2026

Commits on May 4, 2026

Add sampling backend abstraction
tdene
committed
ee3d7c5
Add FlashInfer sampling backend
tdene
committed
a225d8c
Cache-key sampling on padded batch dimensions
tdene
committed
c5ebb27
Capture sampling graphs during warmup
tdene
committed
23b1bf2
Add tests
tdene
committed
a2c354f
Self-review: misc cleanups
tdene
committed
7d31532
Fix: sample_kernel must return its output tensor
tdene
committed
1d6623e
Remove unnecessary tensor copy
tdene
committed
2ec02d7
Fix speculative graphing
tdene
committed
fd31b12
lint
tdene
committed
6b274c4
Remove unnecessary D2D copy
tdene
committed
e089240
Clean up after merge
tdene
committed
bbfa111
Expand test coverage to FlashInfer
tdene
committed
ba2566f
Self-review
tdene
committed
ed80ad9
Default to FlashInfer sampling
tdene
committed
60955a2
Fix torch sampling unit tests
tdene
committed
0c3666c
Fix topk & topp tests
tdene
committed
d6affe7
lint, lint, lint away
tdene
committed
583f1d2
fail the user
tdene
committed
a10e725
Clamp temperature in FlashInfer sampling

tdene
and
santhnm2
committed
92b5039
Cover temp 0 in test_sample_from_dynamic_logits

tdene
and
santhnm2
committed
00829e8
Resolve the rest of the rebase conflict
tdene
committed
de5c56d
Revert the default to torch sampling
tdene
committed
e580bde