Skip to content

fix(kernel): make Triton top-k/top-p sampling exact - #329

Open
taking-lying-flat wants to merge 3 commits into
FlashML-org:mainfrom
taking-lying-flat:fix/triton-topk-correctness
Open

fix(kernel): make Triton top-k/top-p sampling exact#329
taking-lying-flat wants to merge 3 commits into
FlashML-org:mainfrom
taking-lying-flat:fix/triton-topk-correctness

Conversation

@taking-lying-flat

@taking-lying-flat taking-lying-flat commented Sep 1, 2026

Copy link
Copy Markdown

Summary

  • incorporate the exact single-launch fp32 histogram search from fix(kernel): exact single-launch triton top-k/top-p sampling #345
  • clip top-k and top-p boundary ties deterministically by token id, including combined top-k → top-p filtering
  • retry the same exact kernel with one CTA per row when cooperative launch is unavailable
  • scope cooperative-launch capability by device and operation, and propagate unrelated Triton failures
  • remove the older approximate multi-launch fallback and redundant workspace/target conversions

This keeps the final PR to one production file; no test files are included.

Validation

  • local-only CUDA regression suite: 9 passed
  • randomized 64 × 4096 comparison against stable PyTorch top-k/top-p references, including sampled-token support checks
  • Ruff, Python byte-compilation, and diff whitespace checks
  • RTX A1000 benchmark against fix(kernel): exact single-launch triton top-k/top-p sampling #345: comparable latency, with no additional full-vocabulary tie pass

Credit

This revision builds directly on #345. Thank you, @jason-fxz, for taking the original correctness report seriously, replacing the fragile estimator with a much stronger exact design, and carefully reviewing #329. The core single-launch histogram approach here is your work; these follow-up changes tighten its boundary semantics and fallback behavior. I genuinely appreciate the time and thought you put into it.

@taking-lying-flat taking-lying-flat changed the title fix(sampling): handle Triton top-k candidate limits fix(kernel): make Triton top-k/top-p sampling exact Sep 2, 2026
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.

2 participants