Conversation
transform_index_page_table_{decode,prefill}_fast both assert the top-k
width is exactly 2048. k-pool appends up to index_kpool - 1 open-tail
tokens to index_topk, so GLM-5.3-Flash hands over 2051 and both entry
points reject it.
Only the decode kernel actually needs the width. It bakes TOPK in as a
literal constexpr, reuses it as the row stride, and covers a row with one
unmasked tl.arange(0, TOPK) -- which requires a power of two. Add a tiled
variant that carries the strides explicitly and masks the trailing
partial tile, and keep the 2048 single-program path for the unpooled
width so nothing regresses there.
The prefill kernel is already width-generic: it takes TOPK as a kernel
argument, masks on topk_offsets < TOPK, and its wrapper already sizes the
grid with cdiv(topk_indices.shape[1], block_topk). Its assert only
rejected widths the kernel handles, so drop it.
Parametrize the existing test on the top-k width and add 2051 (the
GLM-5.3-Flash k-pool width), 515 (a partial trailing tile) and 257 (below
one tile) for both entry points. Register it on the AMD mi35x per-PR
suite alongside the existing CUDA registration -- the kernel is shared
and the widths are not ROCm-specific.
Co-authored-by: Cursor <cursoragent@cursor.com>
Jacob0226
marked this pull request as ready for review
September 2, 2026 10:30
Jacob0226
requested review from
BBuf,
DarkSharpness,
HaiShaw,
HydraQYH,
celve and
yuan-luo
as code owners
September 2, 2026 10:30
5 tasks
Contributor
Author
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
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.
Summary
This is the unfused DSA top-k path, the counterpart to the fused one in #37563. Both modes run that kernel, so this is a fallback fix, not a substitute for it.
SGLANG_DSA_FUSE_TOPKtruefast_kpool_topk_transform_fuseddoes the top-k and the page-table gather, so this PR's entry points are never reachedfalsetransform_index_page_table_{decode,prefill}_fastgathers--dsa-topk-backend torch; PD disaggregation with MTP index share, whereshould_use_dsa_fused_topk()returns false whatever the env says; or set by hand to bisectBoth of those entry points assert the top-k width is exactly 2048. k-pool appends up to
index_kpool - 1open-tail tokens toindex_topk, so GLM-5.3-Flash hands over 2051 and both reject it.Scope
transform_index.pyonly. No environment variable, no architecture gate.transform_index_page_table_decode_tiled_kernelAssertionErrorunder the same configurationTest plan
Docker:
rocm/sgl-dev:v0.5.18-rocm720-mi35x-20260901.The two 2048 asserts and the tiled kernel
transform_index_page_table_decode_tiled_kernelare the only variables.cd test && python3 registered/kernels/ops/attention/test_dsa_transform_index.py.AssertionErrorWidths are compared against
torch.gatherover the page table with-1for unselected slots, the reference the file already used.register_amd_ci(suite="stage-b-test-1-gpu-small-amd-mi35x")added beside the file's existingregister_cuda_ci. The kernel is shared and the widths come from the model config, so both platforms want the coveragepre-commitvalidate registered test CI registriesCI States
Latest PR Test (Base): ❌ Run #33619810158
Latest PR Test (Extra): ❌ Run #33628232668
Latest PR Test (AMD ROCm 7.2): ❌ Run #33619810064