Fix DSV4 C4 top-k v2 dual output in one pass - #38016
Closed
AMD-yanfeiwang wants to merge 3 commits into
Closed
AMD-yanfeiwang wants to merge 3 commits into
AMD-yanfeiwang wants to merge 3 commits into
Conversation
AMD-yanfeiwang
requested review from
BBuf,
DarkSharpness,
Fridge003,
HaiShaw,
HydraQYH,
Qiaolin-Yu,
celve,
hebiao064,
ispobock,
merrymercy and
yuan-luo
as code owners
September 4, 2026 13:56
Collaborator
|
I think this is a better implementation. Let's move forward with it. |
Collaborator
|
maybe duplicate of #33672 |
Contributor
Author
|
@DarkSharpness Agreed, thanks for pointing this out. Looking forward to #33672 landing. Once it is merged, I’ll close this PR and follow up with a separate PR for the ROCm-specific optimizations and additional test coverage from this work. |
Contributor
Author
|
Solved by #33672 |
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
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
Fixes #37892.
DSV4 paged prefill produces two representations of the same C4 selection:
The generic DSV4 backend allocates
c4_sparse_raw_indicesfor paged prefill, but the C4 dispatch only selected top-k v2 whenraw_indices is None. Those prefills therefore fell back to the legacy v1 transform.ROCm's ordinary non-CP HIP backend does not request this raw buffer, but CP prefill, indexer capture, and HiSparse decode can. #36684 explicitly enabled top-k v2 for DeepSeek-V4 on HIP, yet each of those raw-output paths still falls back to the AOT v1 operator. The exact GB300 crash is architecture-specific; the raw-output dispatch and performance issue is shared.
#35041 removed v2's old optional raw output because no production caller used
it and page-only calls paid a nullable-pointer branch. Paged prefill now is that production caller.
This takes the same dispatch direction as #37991, but emits both outputs in one v2 invocation rather than repeating the complete top-k selection.
Changes
PAGE_TABLE_AND_INDICESv2 mode. It writes aligned raw and paged outputs from the same selected indices.sgl-kernelbackend to v2 even when raw indices are requested.Correctness
MI355X, ROCm 7.2, PyTorch 2.11:
test_topk_v2.py: 297 passedTwo independent v2 calls select the same sets but do not preserve the same output ordering between the raw and paged calls. The fused mode derives both buffers from one selection and keeps them slot-aligned.
Performance
MI355X, ROCm 7.2, PyTorch 2.11, median of 8 interleaved samples.
Lis the C4 sequence length.The fused mode is within -0.6% to +2.7% of page-only v2 across those shapes, so the second output is nearly free compared with repeating selection.
DeepSeek-V4-Pro TP8/DCP8 (which requests raw C4 indices), chunked prefill 16384, concurrency 2, one output token, five repeats:
There is no measurable end-to-end improvement in this configuration; these deltas are within system noise. The kernel-level result is the performance claim.
Testing
Pre-commit checks passed for isort, ruff, ruff-format, codespell, and clang-format.
CI States
Latest PR Test (Base): ❌ Run #33961904069
Latest PR Test (Extra): ❌ Run #33961903803
Latest PR Test (AMD ROCm 7.2): ❌ Run #33961903978