fix(dsa): correct packed FlashInfer top-k and backend selection semantics - #32490
Merged
Fridge003 merged 3 commits intoJul 31, 2026
Merged
Conversation
zianglih
requested review from
1am9trash,
Fridge003,
HaiShaw,
Qiaolin-Yu,
YAMY1234,
hubertlu-tw,
ispobock,
kkHuang-amd,
merrymercy and
rainj-me
as code owners
July 27, 2026 06:55
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
zianglih
changed the title
Contributor
Author
|
/tag-and-rerun-ci |
zianglih
changed the title
This was referenced Jul 27, 2026
Contributor
Author
|
/rerun-failed-ci |
1 similar comment
Contributor
Author
|
/rerun-failed-ci |
ziang-and
force-pushed
the
fix/flashinfer-packed-paged-dsa-topk
branch
from
July 28, 2026 22:08
529b7aa to
a6a14aa
Compare
Packed PAGED extend uses batch-global score offsets with request-local page tables, which FlashInfer's single row_starts argument cannot represent. Fall back to the SGL fused transform for this path and cover shifted multi-query batches in the backend equivalence test. Co-authored-by: Parth Chadha <parth@humansand.ai>
Pass the known logit row count to dynamic repeat_interleave calls so PyTorch does not synchronize CUDA to infer the output size. Extend the existing fused backend-equivalence coverage to assert the no-sync call contract. Co-authored-by: Parth Chadha <parth@humansand.ai>
ziang-and
force-pushed
the
fix/flashinfer-packed-paged-dsa-topk
branch
from
July 28, 2026 22:10
a6a14aa to
70a9ba0
Compare
aleozlx
pushed a commit
to flashinfer-ai/flashinfer
that referenced
this pull request
Jul 30, 2026
<!-- .github/pull_request_template.md --> ## 📌 Description @HumansAnd SGLang needs `top_k_page_table_transform` to support independent score-window and page-table starts in its fused packed PAGED DSA path: - `row_starts` identifies the score window used for top-k selection. - `page_table_row_starts` identifies the page-table window used to translate the selected local indices. The current FlashInfer API applies `row_starts` to both operations, so it cannot represent this case. SGLang must therefore fall back to the SGL kernel even when `--dsa-topk-backend flashinfer` is selected. That produces correct indices, but it bypasses FlashInfer's deterministic, tie-break, and graph-safe fused top-k behavior for this path. This PR adds optional `page_table_row_starts` support. When it is omitted, FlashInfer continues to use `row_starts` for both operations, preserving the existing API behavior. The separate start is propagated through the Python and trace APIs, TVM FFI binding, radix and filtered implementations, deterministic post-sort, graph-safe dispatch, and the trivial `length <= k` path. Once SGLang adopts a FlashInfer release containing this API, it can remove the packed PAGED fallback and use FlashInfer fused top-k with the intended backend semantics. This PR does not change the SGLang call site. ### API Design For each output row `i`, define: ```text batch_i = row_to_batch[i] if row_to_batch is provided, else i score_start_i = row_starts[i] if row_starts is provided, else 0 page_start_i = page_table_row_starts[i] if page_table_row_starts is provided, else score_start_i length_i = lengths[i] ``` Top-k selection produces local offsets `local_idx[i, j]` in `[0, length_i)` by ranking: ```text input[i, score_start_i + local_idx[i, j]] ``` The page-table transform uses the same local offsets but an independent page-table origin: ```text output[i, j] = src_page_table[batch_i, page_start_i + local_idx[i, j]] ``` **`page_table_row_starts` affects only the page-table lookup used to produce output values.** It does not change the input score window, the selected local offsets, or the output tensor shape; it only changes which columns of `src_page_table` are gathered into the output. The argument responsibilities are therefore orthogonal: - `row_to_batch` selects only the row of `src_page_table`; multiple score rows may map to the same page-table row. - `row_starts` selects only the score-window origin. - `page_table_row_starts` selects only the page-table-window origin. It may be provided independently of `row_starts`; when omitted, it reuses `row_starts` for backward compatibility. If `length_i <= k`, all local offsets `0..length_i-1` are transformed and the remaining output positions are `-1`. Otherwise, exactly `k` local offsets are selected according to the existing deterministic and tie-break semantics. Each optional mapping/start tensor has shape `(num_rows,)`, dtype `int32`, and resides on the same CUDA device as `input`. Callers must satisfy `0 <= length_i`, `0 <= batch_i < src_page_table.shape[0]`, `0 <= score_start_i`, `score_start_i + length_i <= input.shape[1]`, `0 <= page_start_i`, and `page_start_i + length_i <= src_page_table.shape[1]`. This selection-plus-gather contract is not specific to SGLang: it represents any packed layout where score storage and lookup-table storage use different origins. Absolute starts are used instead of deltas so the API does not assume a relationship between the two windows or expose framework-specific metadata such as `cu_seqlens`. ## 🔍 Related Issues - SGLang DSA top-k backend integration: sgl-project/sglang#22851 - SGLang packed PAGED correctness fallback and backend-selection cleanup: sgl-project/sglang#32490 ## 🚀 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 on an NVIDIA B200 using the editable source build and source JIT: ```text pre-commit run --all-files Passed python3 -m pytest -vv -s tests/utils/test_topk.py -k 'test_top_k_transform_with_row_starts' 48 passed, 1334 deselected, 2 warnings in 0.61s ``` The test extends the existing `test_top_k_transform_with_row_starts` Cartesian product across radix/filtered dispatch, graph-safe mode, deterministic mode, shared/separate starts, and both trivial and selected rows. ## Reviewer Notes Review focus is welcome on propagation through the deterministic post-sort and graph-safe filtered paths, where page-table translation occurs separately from score selection. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added optional `page_table_row_starts` support to the fused Top‑K page-table transform for independent per-row destination window offsets. * Updated tracing and reference implementations to model separate score-window (`row_starts`) and destination page-table-window (`page_table_row_starts`) offsets. * **Bug Fixes** * Corrected page-table addressing when the score and destination windows start at different offsets. * **Tests** * Expanded Top‑K transform coverage to include deterministic mode and separate `page_table_row_starts`, with additional trace-based reference checks. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
Fridge003
approved these changes
Jul 30, 2026
Collaborator
|
/rerun-test test/registered/models_e2e/test_dsa_glm52_tp_mtp.py test/registered/models_e2e/test_dsa_glm52_dp_mtp.py test/registered/models_e2e/test_dsa_glm52_nvfp4_dp_mtp.py test/registered/models_e2e/test_dsa_glm52_nvfp4_tp_mtp.py test/registered/cp/test_glm52_cp_index_share.py test/registered/models_e2e/test_dsa_glm52_hisparse.py test/registered/models_e2e/test_dsa_glm52_cache_layer_split.py test/registered/kernels/ops/attention/test_dsa_indexer.py |
Contributor
|
Results for 🚀 🚀 🚀 🚀 |
5 tasks
bkryu
pushed a commit
to flashinfer-ai/flashinfer
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 -->
saturn-acc
pushed a commit
to saturn-acc/sglang
that referenced
this pull request
Aug 16, 2026
…tics (sgl-project#32490) Co-authored-by: Parth Chadha <parth@humansand.ai>
5 tasks
jakki-amd
pushed a commit
to jakki-amd/sglang
that referenced
this pull request
Sep 9, 2026
…tics (sgl-project#32490) Co-authored-by: Parth Chadha <parth@humansand.ai>
Atituiset
pushed a commit
to Atituiset/sglang
that referenced
this pull request
Sep 10, 2026
…tics (sgl-project#32490) Co-authored-by: Parth Chadha <parth@humansand.ai>
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
2 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.
Motivation
@HumansAnd
PR #22851 introduced
--dsa-topk-backendand integrated FlashInfer and PyTorch top-k into the fused and unfused DSA paths. This PR follows up on that integration with three focused fixes:row_startsargument to both score selection and page-table lookup, so it cannot represent this layout.SGLANG_OPT_USE_TOPK_V2is ansgl-kernelimplementation optimization and should not override an explicit--dsa-topk-backendchoice. Previously, eligible decode calls could enter the SGL top-k v2 path even when another backend, such asflashinferortorch, was selected. CUDA graph metadata construction made the same assumption when deciding whether to build the v2 plan and retain the page-size-1 table.torch.repeat_interleavecalls withoutoutput_sizeinfer their output size by reading repeat counts back from CUDA. Each affected mapping is already known to produce one entry per logit row, so supplying that row count avoids the host synchronization without changing mapping semantics.The accompanying page-size lookup change does not introduce new paging behavior.
_topk_transform_v2_pagedalready consumesDSAMetadata.real_page_table, and both eager and CUDA graph metadata associate that table withDSAMetadata.page_sizefrom the backend'sreal_page_size. Reading the table and its page size from the same metadata removes an ambient forward-context dependency while preserving the existing value.Modifications
sgl_kernel.fast_topk_transform_fusedas a correctness fallback for the unsupported FlashInfer PAGED case whererow_startsis present. Unshifted PAGED and RAGGED calls continue to use FlashInfer.SGLANG_OPT_USE_TOPK_V2subordinate to--dsa-topk-backend: top-k v2 is eligible only whensgl-kernelis selected. Apply this policy consistently to dispatch, plan construction, and CUDA graph wide-page-table elision.DSAMetadatathat suppliesreal_page_table, rather than consulting global forward context.output_sizeto dynamic repeat-interleave calls in FlashInfer PAGED row mapping and top-k offset construction.Known Limitation
The packed PAGED fallback in this PR is numerically correct, but it is not a complete FlashInfer fused top-k path: when
row_startsis present, the FlashInfer branch invokessgl_kernel.fast_topk_transform_fused.A fully semantically correct FlashInfer fused top-k implementation for this case requires flashinfer-ai/flashinfer#4169. That PR adds a separate
page_table_row_startsargument so FlashInfer can use the score-window and page-table starts independently. Until SGLang adopts a FlashInfer version containing that API,--dsa-topk-backend flashinferwill continue to use the SGL fallback for packed PAGED rows, and the FlashInfer deterministic and tie-break settings will not apply to that shape.The backend-selection fix in this PR is independent of that dependency: it prevents the SGL-only top-k v2 optimization from overriding an explicitly selected non-SGL backend on paths the selected backend supports.
Accuracy Tests
Focused tests were run on an NVIDIA B200 with CUDA 13.0, FlashInfer 0.6.15.post1, and
sglang-kernel0.4.5:Result:
3 passed, 30 subtests passed in 11.58s.The tests compare exact selected index sets using tie-free logits across PAGED and RAGGED layouts, shifted and unshifted rows, and multi-query batches. The multi-query fused cases also assert that every dynamic top-k repeat supplies
output_size=num_rows, the PyTorch API contract that avoids readingsum(repeats)back to the host. FlashInfer tie-break modesNone,small, andlargeare exercised on paths that invoke FlashInfer. Shifted PAGED cases validate the SGL fallback against the SGL backend and therefore do not exercise FlashInfer tie-breaking. No model-level accuracy benchmark was run.pre-commit run --all-filesalso passed.Speed Tests and Profiling
No standalone performance benchmark was run for the repeat-interleave change. Supplying the already-known
output_sizeremoves PyTorch's implicit host read without changing mapping values. The packed PAGED change is a correctness fix, while the top-k v2 change restores the expected backend-selection semantics. FlashInfer remains unchanged for layouts its API can represent; only packed PAGED rows with separate score/page-table coordinate systems use the existing SGL fused transform.Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): 🚫 Run #30403592750
Latest PR Test (Extra): ❌ Run #30403592267