fix(kernel): make Triton top-k/top-p sampling exact - #329
Conversation
|
Thanks for folding #345 in and for the follow-up work — the per-device I'm planning to merge this PR instead of #345. One change before that: the tie clipping in |
|
LGTM. |
Traz o commit oficial novo: - fix(kernel): make Triton top-k/top-p sampling exact (FlashML-org#329) Sem conflito com o trabalho local do fork -- confirmado via 'git merge-tree --write-tree', arquivos tocados nao se sobrepoem.
…lashML-org#336 hf shard download Cherry-picked from FlashML-org/FreeToken 092ce4a; FlashML-org#343 is redundant with the fork's compressed-tensors NVFP4 detection and FlashML-org#319 conflicts with the kernel backend, both left out. Box tests 408 passed with the known order-dependent parity failure (FlashML-org#18). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A88MCbnLtwsFSHmqwuJezY
Brings in exact Triton top-k/top-p sampling (FlashML-org#329), nvfp4 detection behind mixed-precision compressed-tensors (FlashML-org#343), GLM-5.3-Flash config test, issue templates and README/CONTRIBUTING edits. Trial merge showed no conflicts; README.md and models/config.py were the only overlapping files. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FKGJYhK5m18zMnjm5nD8TE
Brings this fork up to upstream main: exact Triton top-k/top-p sampling (FlashML-org#329), the sm_89 row-wise _scaled_mm workaround (FlashML-org#243), Qwen3.8-Flash-Next and GLM-5.3-Flash, the in-repo fused_topk router, and the safetensors-index shard fix. One conflict, in models/blocks.py, and it is a conflict only by position: upstream adds forward_host_ctx to BaseLLMModel at the same point this fork adds the hidden-state capture DFlash needs. Both are kept -- forward_host_ctx inside BaseLLMModel, the HiddenStateCapture mixin after it -- and both import lines are kept, since each side needed a different one. Verified after resolving: the 32 speculative and detokenizer tests pass, as do 88 scheduler tests and the muse_glimmer weight tests. One order-dependent failure appeared once in test_abort_inflight_prefill and did not reproduce in the same invocation afterwards; it is recorded as an unexplained flake rather than declared fixed.
Summary
Dropping the token-id tie clipping also removes its per-CTA tie workspace and one synchronization stage from the fused path. This keeps the final PR to one production file; no test files are included.
Validation
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; the follow-up changes retain its exactness while making device handling and fallback behavior more robust.