[FlashInfer V0.6.18] feat(dsv4): support --dsa-topk-backend flashinfer with fused top-k - #33237
Merged
Fridge003 merged 7 commits intoSep 1, 2026
Merged
Conversation
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
This was referenced Aug 2, 2026
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 -->
ziang-and
force-pushed
the
agent/dsv4-flashinfer-fused-topk
branch
from
August 9, 2026 06:17
0534ef4 to
fd150ca
Compare
zianglih
changed the title
zianglih
changed the title
zianglih
marked this pull request as ready for review
August 10, 2026 07:11
zianglih
requested review from
1am9trash,
CatherineSue,
Fridge003,
HaiShaw,
JustinTong0323,
Qiaolin-Yu,
YAMY1234,
hebiao064,
hubertlu-tw,
ispobock,
kkHuang-amd,
merrymercy,
rainj-me and
slin1237
as code owners
August 10, 2026 07:11
41 tasks
ziang-and
force-pushed
the
agent/dsv4-flashinfer-fused-topk
branch
from
August 16, 2026 06:11
4b99b82 to
1e15e64
Compare
zianglih
changed the title
Contributor
Author
|
rebasing, now depending on #36954 |
ziang-and
force-pushed
the
agent/dsv4-flashinfer-fused-topk
branch
from
August 29, 2026 06:57
798019b to
08749f3
Compare
ziang-and
force-pushed
the
agent/dsv4-flashinfer-fused-topk
branch
from
August 29, 2026 07:00
08749f3 to
d05fef2
Compare
mmangkad
reviewed
Aug 29, 2026
mmangkad
left a comment
mmangkad
left a comment
Collaborator
There was a problem hiding this comment.
Just lint to fix I think, plus a few small inline comments. Ran DSV4 on GB300 with 0.6.18, looks good.
mmangkad
approved these changes
Aug 30, 2026
Contributor
Author
|
/tag-and-rerun-ci |
Contributor
Author
|
/tag-and-rerun-ci |
Contributor
Author
|
nv ci green |
Fridge003
approved these changes
Sep 1, 2026
1am9trash
added a commit
to RolaoDenthu/sglang
that referenced
this pull request
Sep 1, 2026
This was referenced Sep 1, 2026
This was referenced Sep 1, 2026
StevenChenSE
pushed a commit
to StevenChenSE/sglang
that referenced
this pull request
Sep 6, 2026
…r with fused top-k (sgl-project#33237)
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
3 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
Scope: DeepSeek V4 FlashInfer top-k backend support. This PR makes
--dsa-topk-backend flashinferusable for the DeepSeek V4 indexer. It uses FlashInfer's fused compact-page transform whenSGLANG_DSA_FUSE_TOPK=trueand preserves the existing unfused FlashInfer path when the flag is disabled.The packed PAGED-row prerequisite was merged in #33006. DeepSeek V4 adds a compact page table where one page-table entry represents 64 score positions. FlashInfer #4315 added the required
page_size, caller-owned output, raw-index output, and padded-stride support.Modifications
top_k_page_table_transformwith:page_size=64;flashinfer.top_kplus vectorized page translation path and resolve the fused/unfused transform once during backend initialization.should_use_topk_v2()decision into indexer metadata; FlashInfer and Torch do not create or copy an unused SGL plan.flashinfer-python>=0.6.18when FlashInfer is selected for target or speculative-draft DSA top-k, while preserving the existing resolved attention-backend check.bench_topk.pyto benchmarkpage_size={1,64}as a Cartesian axis with a preallocated output while keeping the focused CI value atK=512. The expanded performance grid below uses a source-unchanged runtime-axis wrapper. Page size 64 matches the compact CUDA DSA layout; page size 1 is the noncompact baseline.K={512,1024}, optional raw output, padded score strides, replayed inputs, short rows, and raw/translated alignment. Tests exercise non-contiguous page tables and assert FlashInfer 0.6.18's-1padding contract for every unused translated/raw suffix slot.Accuracy Tests
Post-merge validation used SGLang
abaeebc57eacb256a4e8013b59646c4d71cf0dc6on one NVIDIA B300:d8750c398dccbea72bbc37bf37ddacae2defe801and upstreammain881cbfe54c98356cfa1eaa134aa4d0be702fc90f.lmsysorg/sglang:nightly-dev-cu13-20260831-bb5e6198(amd64 imagesha256:3c75c6f7f8e6eaeebdcf92ef47caef818608134252ad5898f4c670b7911768a6).2.13.0+cu130, CUDA 13.0.flashinfer-python==0.6.18,flashinfer-cubin==0.6.18, andflashinfer-jit-cache==0.6.18+cu130, force-reinstalled directly from the final release assets because this nightly image predates [Deps] Bump FlashInfer to 0.6.18 #36954.Install and focused test commands:
Raw results:
The suite covers upstream's pre-publish/platform guards, init-time fused/unfused routing, required resolved top-k-v2 metadata, a non-contiguous compact page table,
K={512,1024}, optional raw output, padded score strides,-1suffix padding, replayed inputs, short rows, raw/translated alignment, and CUDA-graph replay. A direct mock smoke independently checked both target/speculative FlashInfer version gates and the caller-resolved XPU exclusion. Changed-file pre-commit passed and left the checkout clean. No DeepSeek V4 model-serving accuracy run or performance rerun was performed; the speed data below remains historical pre-merge evidence. Test log SHA-256:341cb33a65a9e7f97a6c9985a6977a1d111074d66960040146776aaf3714ff10.Speed Tests and Profiling
Historical single-GPU operator microbenchmark at SGLang
798019b5dcf815371b4986b84117ad56eff092ae, using the pre-release FlashInfer 0.6.18 nightly recorded below:K={512,1024,2048} × seq_len={8192,65536,131072} × batch_size={4,8,16,32,64} × page_size={1,64}. Providers: SGL JIT-v1, SGL JIT-v2, FlashInfer, and PyTorch. JIT-v1 is unavailable atK=2048by design.K=512. A source-unchanged inline wrapper accepts the four requested axes as command-line arguments and replaces the benchmark's runtime configuration before callingBenchmark.run(); the provider list is asserted unchanged.FLASHINFER_TOPK_ALGOunset (automatic);SGLANG_DSA_FUSE_TOPKunset. The benchmark directly calls the fused FlashInfer API, so this is not a fused-versus-unfused comparison. The tested FlashInfer nightly predates the open CUB TopK PR, so these are not CUB-backend results. Each process emitted the already-documented AOT-to-JIT fallback warning; compilation remained outside timing.Command:
Complete raw latency output:
Cross-sweep medians;
FI / JIT-v2is FlashInfer latency divided by JIT-v2 latency:JIT-v2 won every same-page shape. FlashInfer/JIT-v2 geometric means for page sizes 1/64 were 3.28x/3.30x at
K=512, 3.24x/3.28x atK=1024, and 3.07x/3.15x atK=2048; the overall same-page geometric mean was 3.22x. Page size 64 reduced latency relative to page size 1 by 3.30% for JIT-v2 and 1.88% for FlashInfer. FlashInfer remained 2.83x faster than the PyTorch fallback by geometric mean.For the DeepSeek V4 routing proxy (
K=512, page_size=64), FlashInfer used 3.30x JIT-v2 latency across the 15 shapes. Per-sequence geometric means were 2.42x at 8,192, 3.35x at 65,536, and 4.43x at 131,072. At batch 16, the exact ratios were 12.1856/5.0381 us = 2.42x, 47.5974/16.6707 us = 2.86x, and 77.8629/15.5548 us = 5.01x.For the current GLM-5 routing proxy (
K=2048, FlashInfer page size 1 versus JIT-v2 page size 64), FlashInfer used 3.23x JIT-v2 latency across the 15 paired shapes (range 2.32x-6.55x). Per-sequence geometric means were 2.39x at 8,192, 3.15x at 65,536, and 4.48x at 131,072. At batch 16, the exact ratios were 12.9030/5.3658 us = 2.40x, 51.9648/20.3158 us = 2.56x, and 88.1805/17.0515 us = 5.17x.These are CUDA-graph operator results, not serving results. Inputs use random FP32 scores, full lengths, identity page tables, contiguous tensors, and preallocated outputs; the grid does not cover ties, short rows, padded score strides, or raw-index output. PyTorch uses sorted top-k. The 131,072-token case is substantially closer to GLM-5 than the earlier grid but remains below its 202,752-token maximum. No DeepSeek V4 or GLM-5 end-to-end throughput result is claimed.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ✅ Run #33354611262
Latest PR Test (Extra): ❌ Run #33354611056
Latest PR Test (AMD ROCm 7.2): ❌ Run #33354611147