perf(topk): skip output index sort for tie-break selection - #4295
Merged
Merged
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📥 CommitsReviewing files that changed from the base of the PR and between 89c649679af75792b2cad1020c91b41bf2262adc and 138b943. 📒 Files selected for processing (5)
💤 Files with no reviewable changes (1)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughTop-K tie-break handling no longer implies deterministic output ordering. API dispatch preserves caller-selected modes, clustered execution rejects explicit tie-break requests, and configurable finalization handles sorting and transformed outputs. Benchmarks and tests cover the revised behavior. ChangesTop-K tie-break behavior
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant TopKAPI
participant CUDADispatch
participant LaunchFinalizeTopKIndices
Caller->>TopKAPI: request top-k with deterministic and tie_break
TopKAPI->>CUDADispatch: pass both modes independently
CUDADispatch->>LaunchFinalizeTopKIndices: pass index-sorting mode
LaunchFinalizeTopKIndices-->>Caller: return plain or transformed output
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
zianglih
marked this pull request as ready for review
July 31, 2026 10:08
zianglih
requested review from
aleozlx,
bkryu,
cyx-6,
jimmyzho,
kahyunnam,
nv-yunzheq,
saltyminty,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
July 31, 2026 10:08
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
zianglih
requested review from
Aneureka,
dhiraj113,
feih-nv,
jiahanc,
qiching,
qsang-nv and
samuellees
as code owners
July 31, 2026 10:08
69 tasks
yzh119
reviewed
Jul 31, 2026
yzh119
left a comment
yzh119
left a comment
Collaborator
There was a problem hiding this comment.
It should be an important feature to have, thanks for the contribution.
| @@ -24,6 +24,8 @@ | |||
| from flashinfer.testing.utils import bench_gpu_time | |||
| from flashinfer.utils import get_compute_capability | |||
|
|
|||
| BENCH_ENABLE_CUPTI = True | |||
Collaborator
There was a problem hiding this comment.
How about reading this value from environment variable?
Contributor
Author
There was a problem hiding this comment.
I have dropped this interface. The instability is now avoided by using enable_cupti=True + use_cuda_graph=False, which is already a common pattern in the repo.
Collaborator
|
/bot run |
Collaborator
5 tasks
Collaborator
|
[FAILED] Pipeline #60457495 — 6/18 executed test jobs passed Compared with nightly #60449790. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsPre-existing failures
Timeouts, infrastructure, or incomplete jobs
|
zianglih
force-pushed
the
agent/topk-tie-break-unsorted
branch
from
August 2, 2026 17:59
89c6496 to
138b943
Compare
Collaborator
|
/bot run |
Collaborator
Collaborator
|
[FAILED] Pipeline #60826663 — 15/18 executed test jobs passed Compared with nightly #60712014. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsPR-related regressions
New relative to nightly (attribution uncertain)
Pre-existing failures
Timeouts, infrastructure, or incomplete jobs
|
IwakuraRein
approved these changes
Aug 4, 2026
5 tasks
bkryu
approved these changes
Aug 5, 2026
Collaborator
There was a problem hiding this comment.
Thank you @zianglih.
Key change here is that previously if we set tie_break to small or large and set deterministic=False, the outputs were still deterministic and sorted because we silently sorted. This was not a guaranteed behavior
This PR seems to make the output not sorted so the non-promised behavior is changing. Non-blocking but wanted to note in this PR. Will trigger the CI one last time to get targeted testing and approve.
Collaborator
|
/bot run tests/utils |
Collaborator
Collaborator
|
[CANCELING] Pipeline #61240246: canceled |
bkryu
pushed a commit
that referenced
this pull request
Aug 9, 2026
<!-- .github/pull_request_template.md --> ## 📌 Description @HumansAnd SGLang's DeepSeek V4 indexer uses a compact page table in which one entry represents 64 score positions. Its score rows may also have padding between rows, and CUDA graph capture owns the translated and raw-index output buffers. The current `top_k_page_table_transform` contract assumes one page-table entry per score and allocates only the translated output, so SGLang has to split this into `top_k`, score re-gathering, compact-page translation, and output copies. This PR extends the existing fused page-table transform for that layout: - Add `page_size`, defaulting to `1` for backward compatibility. - Add optional caller-owned `out` and `out_raw_indices` buffers. Raw indices remain positionally aligned with translated indices, including deterministic post-sort and `-1` padding; the two buffers must be disjoint. - Honor the input row stride instead of requiring tightly packed rows, while preserving last-dimension contiguity and alignment-safe vectorization. - Propagate the contract through the Python API, TVM FFI binding, Radix and Filtered implementations, graph-safe dispatch, deterministic post-sort, and the trivial `length <= k` path. - Extend only the operation-local trace template for `page_size`. Its core one-output definition excludes destination buffers; raw-output calls are deliberately not emitted and use a distinct routing identity so Trace Apply falls back to the API. For each selected local score index `idx`, the compact transform is: ```text physical_page = src_page_table[ batch_idx, page_table_row_start + idx // page_size ] output = physical_page * page_size + idx % page_size ``` `page_table_row_starts` is measured in page-table entries, while `row_starts` is measured in score elements. When `page_size > 1` and `row_starts` is supplied, callers must therefore provide `page_table_row_starts` explicitly rather than relying on the existing shared-start behavior. The C++ path uses policy types rather than a second boolean mode. The FFI boundary selects `DirectPageTableKernelPolicy` or `ConfigurablePageTableKernelPolicy` once, and the kernel ABI is derived structurally from the policy type: an empty policy contributes zero arguments, while a stateful trivially-copyable policy contributes one object. Host dispatch carries one typed policy value through the selection stack; only the terminal launch forms the zero-or-one argument pack. The direct policy therefore adds no kernel argument or device branch, while the configurable policy owns score-row layout, logical-to-physical translation, and the optional raw-index sink. Future transforms that share the flat row-selection contract can extend the configurable policy without multiplying kernel variants or changing the selection kernels. Direct translation sites retain their original expressions so supported-path SASS is preserved exactly. This is a clean extension of the API introduced in #4169: `page_size=1`, omitted output buffers, and tightly packed inputs retain the existing behavior and cluster fast path. [SGLang #33237](sgl-project/sglang#33237) uses this API to replace its DeepSeek V4 unfused workaround with one graph-safe FlashInfer call. ## ⚡ Performance Fresh performance validation compares the exact rebase base `29196cf437778906c72630dc5d9850de547501de` with head `cf0319a3497e252219544c0a8b4168c6ba598f88` on the same NVIDIA B200 (driver 580.126.09), CUDA 13.2.78, and PyTorch 2.13.0+cu132. Base and head used separate editable source trees and JIT workspaces. `benchmarks/bench_topk.py` is byte-identical on both sides (SHA-256 `57cd1ca61b38120380cb9ea7cf81ae3ee972484724bb2bb785649ae05cc199d9`). The script uses CUPTI (`cupti-python` 13.2.0 and `nvidia-cuda-cupti` 13.2.75), 10 dry runs, 100 measured iterations, cold L2, and the median. After #4295 it already sets `use_cuda_graph=False`, so no temporary benchmark-source edit was needed and the CUPTI plus CUDA graph instability is excluded. Exact PR-body commands: ```bash python3 benchmarks/bench_topk.py \ --op dsa_topk --dtype bf16 --dsa-input-pattern dsa_relu \ --dsa-case all --dsa-topk 2048 --tie-break python3 benchmarks/bench_topk.py \ --op varlen --dtype bf16 --length-dist causal \ --varlen-k 2048 --varlen-q-len 128 --tie-break ``` After #4295 these exact commands keep `deterministic=False`: they measure the default nondeterministic path plus SMALL/LARGE tie selection without the canonical output-order sort. For fair DSA pairing, the confirmation runs used the same command after `torch.manual_seed(1234)` and `torch.cuda.manual_seed_all(1234)`. Current-mode DSA ran base/head/head/base; varlen used base/head/head/base and its built-in per-case seeds. Canonical-output coverage repeated both commands with `--deterministic` on base and head. All comparisons use matched per-case medians; negative PR delta means the head is faster. | workload / mode | default PR delta (worst) | deterministic PR delta (worst) | tie-small PR delta (worst) | tie-large PR delta (worst) | |---|---:|---:|---:|---:| | DSA, current ABBA | `+0.05%` (`+0.13%`) | n/a | `-0.09%` (`+0.08%`) | `-0.01%` (`+0.12%`) | | page-table varlen, current ABBA | `+0.01%` (`+0.10%`) | n/a | `+0.06%` (`+0.19%`) | `+0.05%` (`+0.24%`) | | ragged varlen, current ABBA | `+0.02%` (`+0.08%`) | n/a | `+0.01%` (`+0.13%`) | `-0.02%` (`+0.14%`) | | DSA, explicit deterministic | `+0.03%` (`+0.15%`) | `-0.09%` (`+0.10%`) | `+0.01%` (`+0.41%`) | `-0.09%` (no regressed point) | | page-table varlen, explicit deterministic | `+0.01%` (`+0.10%`) | `-0.18%` (`+0.03%`) | `-0.09%` (no regressed point) | `-0.05%` (`+0.09%`) | | ragged varlen, explicit deterministic | `-0.01%` (`+0.06%`) | `-0.10%` (`+0.01%`) | `-0.01%` (`+0.03%`) | `-0.13%` (`+0.08%`) | All 12 varlen rows had identical `len_min`, `len_mean`, `len_max`, and `triv%` across paired runs. Current-mode suite geomeans are within `+0.06%`, the largest individual delta is `+0.24%`, and deterministic geomeans are flat or faster apart from a `+0.01%` DSA tie-small geomean. This supports no measurable kernel-latency regression after the final rebase. These sweeps exercise the existing direct compatibility paths (`page_size=1`, contiguous rows, no caller-owned outputs). The configured V4 path has no pre-PR API equivalent. Its policy cleanup was separately checked by an eager ABBA host audit using `page_size=64`, padded row stride, raw output, and all three production shapes: its configured 12-case batched end-to-end geomean was `+0.053%`, with a worst point of `+0.380%`. The policy design also has complementary binary evidence from the exhaustive audit performed after #4295: ```text direct: 354/354 affected kernels, 0 normalized SASS/resource/KPARAM mismatches radix=180, filtered=132, finalizer=42 configured: 116 kernels radix=60, filtered=44, finalizer=12 exactly one trailing 24-byte policy object; PageTable mode only ``` The direct variants retained their upstream parameter counts and constant-bank spans. The configured variants add one policy object without a second policy family or boolean template axis. ## 🔍 Related Issues - SGLang integration: sgl-project/sglang#33237 - Independent score/page-table starts: #4169 - SGLang packed-PAGED workaround and backend-selection fix: sgl-project/sglang#32490 - SGLang DeepSeek V4 Top-K backend integration: sgl-project/sglang#31087 ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). Validation used an editable source build on one NVIDIA B200 with CUDA 13.2.78 and PyTorch 2.13.0+cu132. The rebased head was synced to an isolated source tree and JIT workspace on `flashinfer-pr4366-cu132`: ```text pre-commit run --all-files Passed python3 -m pytest -q \ tests/utils/test_topk.py::test_top_k_page_table_transform_misaligned_scores_without_row_starts \ tests/utils/test_topk.py::test_top_k_page_table_transform_compact_pages_cuda_graph_replay 18 passed, 3 warnings in 93.98s python3 -m pytest -q \ tests/topk_varlen/test_topk_varlen.py::test_radix_preallocated_outputs \ tests/topk_varlen/test_topk_varlen.py::test_out_values_ignored_when_return_values_false 6 passed, 39 warnings in 4.41s python3 -m pytest -q \ tests/trace/test_fi_trace_template_consistency.py \ tests/trace/test_template_init.py \ -k top_k_page_table_transform 6 passed, 1 skipped, 972 deselected, 3 warnings in 0.27s python3 -m pytest -q tests/utils/test_topk.py 1495 passed, 3 warnings in 7.56s ``` Final CUDA 13.2 validation was rerun on rebased head `cf0319a3497e252219544c0a8b4168c6ba598f88`. It covers the page-table changes plus the optional-output overlap from #3901 after conflict resolution. `git range-diff`, `git diff --check`, and `pre-commit run --all-files` also passed. The Top-K matrix covers Radix multi-CTA and Filtered dispatch, graph-safe mode, deterministic mode, optional raw output for default and compact page sizes, independent score/page-table starts, compact and default page sizes, padded row strides, misaligned input bases, trivial and selected rows, and CUDA graph replay with mutated inputs. The trace checks cover the operation-local schema and default-argument initialization. A separate smoke check also verified that the Python `page_size <= 2**30` validation matches the native contract. The warnings are existing CUTLASS DSL deprecations from `tests/conftest.py` and `flashinfer/cute_dsl/utils.py`. ## Reviewer Notes Review focus is welcome on the positional pairing of raw and translated outputs across deterministic post-sort, the structural zero-or-one policy ABI, and the page-table transform in the Radix multi-CTA and graph-safe Filtered epilogues. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added support for compact page tables with configurable page sizes. * Added optional raw-index outputs alongside translated physical indices. * Added reusable output buffers for flexible result storage. * Improved support for empty rows, padding, non-contiguous inputs, and physical-page remapping. * **Bug Fixes** * Added validation for page metadata and output compatibility. * Improved CUDA graph and algorithm-path support. * **Documentation** * Updated API and tracing documentation for page sizes, physical indices, and raw-index outputs. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
7 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
@HumansAnd
Follow-up to #3095.
tie_breakdetermines which equal-valued elements are selected at the top-k boundary. Today, requesting a tie-break also promotes the publicdeterministicflag, which runs the index-ordering finalizer even though the default API output is unsorted.This PR separates deterministic selection from deterministic output ordering:
SMALL/LARGEchoose the intended boundary indices;FinalizeTopKIndicesKernelalways finalizes/remaps indices, withSORT_LOCAL_INDICESas the compile-time switch for its optional local-index radix sort;deterministic=Trueretains the existing CUB index sort and repeatable output order;deterministic=Falseskips the finalizer entirely for plain Top-K, while page-table/ragged transforms run only the required remapping work;can_use_clusters_topknow receivestie_breakand rejects tie-break modes at the same early-exit boundary asdsa_graph_safe;sorted=Trueremains unchanged and still requests descending value order.sortedanddeterministicuse different sort keysThese options are independent and should not be conflated:
sorted=Truetorch.topk(..., sorted=True)deterministic=TrueThe deterministic filtered-path index sort is a canonicalization step; it does not imply descending score order. If both options are enabled, FlashInfer first establishes deterministic index order and then applies a stable descending value sort. Equal values therefore retain the deterministic prior order. For page-table and ragged transforms, local indices are canonicalized before remapping, so the returned mapped IDs are not necessarily numerically sorted.
tie_breakis separate from both: it controls which equal-valued elements are selected at the top-k boundary, not the final output order.Performance
Both sides used the same final benchmark driver and ran on the same physical GPU (
CUDA_VISIBLE_DEVICES=2) in theflashinfer-pr4048-cu132devbox from thehellqueue. Speedup isBEFORE / AFTER.sm100)upstream/mainat668a1ba1ca86432c79f6adad37ecfce8d06ec08389c649679af75792b2cad1020c91b41bf2262adcThe shared timer supports CUPTI with and without CUDA graphs. A control sweep with CUPTI + CUDA graphs + cold-L2 reproduced the same illegal-memory-access failure on both upstream
mainand this branch at the first 128K page-table row, so it is not specific to this PR's kernel changes. The final benchmark follows existing FlashInfer CUPTI microbenchmarks by usingenable_cupti=Truewithuse_cuda_graph=False; no benchmark-specific environment, global, or CLI escape hatch remains.Commands:
export CUDA_VISIBLE_DEVICES=2 python benchmarks/bench_topk.py \ --op dsa_topk \ --dtype bf16 \ --dsa-input-pattern dsa_relu \ --dsa-case all \ --dsa-topk 2048 \ --tie-break python benchmarks/bench_topk.py \ --op varlen \ --dtype bf16 \ --length-dist causal \ --varlen-k 2048 \ --varlen-q-len 128 \ --tie-breakSummary (geometric mean across cases):
DSA Top-K results (latency in µs):
decode_b1_q1_l128kdecode_b8_q1_l64kdecode_b32_q1_l128kprefill_b1_q128_l128kVariable-length transform results (latency in µs):
Raw BEFORE output: DSA Top-K
Raw AFTER output: DSA Top-K
Raw BEFORE output: variable-length transforms
Raw AFTER output: variable-length transforms
🔍 Related Issues
tie_breakfor filtered topk #3095.🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used my preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Coverage includes:
SMALL/LARGEtie-break output;deterministic=Trueordering for plain and transform APIs;k=2048.Reviewer Notes
The intended contract is:
tie_breakcontrols deterministic boundary selection only;deterministic=Trueadditionally requests repeatable output ordering;sorted=Truecontinues to request descending value order.The key implementation detail is keeping
DETERMINISTICfor filtered selection. The sharedFinalizeTopKIndicesKernelalways performs final mapping/writeback;SORT_LOCAL_INDICESindependently controls its pre-writeback CUB sort. Transformed tie-break outputs retain the remap path with sorting disabled, while plain tie-break output skips the finalizer entirely.The benchmark driver keeps CUPTI enabled and disables CUDA graph replay, matching other FlashInfer CUPTI microbenchmarks. This avoids adding a benchmark-global variable, environment setting, or CLI fallback solely for this workload.
Summary by CodeRabbit
Behavior Updates
Benchmarking