[AMD] [GLM-5.3-Flash Day 0] Support non-2048 top-k widths in the DSA page-table transform

[Jacob0226]

Summary

GLM-5.3-Flash's sparse attention picks 512 pools, expands them to 2048 tokens, then adds up to 3 leftover tokens from the end of the sequence. So the width it hands downstream is 2051, not 2048 — and the code that receives it asserts the width is exactly 2048, and rejects it.

Nobody has hit this yet: the default setup uses a single combined kernel that does the whole job itself, so the rejecting code never runs. Three supported setups split that job in two, and the second half is the part that rejects. None of them is AMD-specific — the assert fires on CUDA the same way.

Setup	How often it hits
HiSparse	every decode step
PD disaggregation with MTP index share	every prefill, on the draft step
--dsa-topk-backend torch	every step

Scope

transform_index.py only. No environment variable, no architecture gate. The 2048 path keeps its own kernel, so every current caller keeps its performance:

Width	Decode	Prefill
2048	existing single-program kernel, untouched	untouched
anything else	new tiled kernel, trailing tile masked	already width-generic; only its assert is dropped

Only the decode kernel needed work. It bakes 2048 in as a literal constexpr, reuses it as the row stride, and covers a row with one unmasked tl.arange, which requires a power of two. The prefill kernel already takes the width as an argument and masks on it, so its assert only rejected widths the kernel handles.

The tiled kernel produces identical output to the 2048 one at that width. The branch is there to leave the existing path alone, not because the tiled kernel is wrong.

Test plan

Unit base rocm/sgl-dev:v0.5.18-rocm720-mi35x-20260901. The two 2048 asserts and the tiled kernel transform_index_page_table_decode_tiled_kernel are the only variables. cd test && python3 registered/kernels/ops/attention/test_dsa_transform_index.py.

Build	Result
Baseline	6 failed, 3 passed — every non-2048 case raises AssertionError
This PR	9 passed on gfx950

Widths are compared against torch.gather over the page table with -1 for unselected slots, the reference the file already used.

End to end, since no single PR here can serve the model on its own: main plus the eight open Day-0 PRs — #39338, #39339, #39340, #39341, #38545, #38546, #38547 and #39317 — which all merge cleanly onto it. rocm/sgl-dev:v0.5.19-rocm720-mi35x-20260909, MI355X TP4, decode CUDA graphs on, 1,319 GSM8K examples, one scoring pass each. Only the checkpoint differs between the two rows:

Checkpoint	GSM8K
zai-org/GLM-5.3-Flash, block FP8	97.27%
amd/GLM-5.3-Flash-Quark-MXFP4	96.51%

Accuracy: not attributable to this PR in isolation, and there is no standalone baseline — GLM-5.3-Flash does not start on ROCm without the rest of the stack.

Check	Detail
Coverage	2051 (GLM-5.3-Flash k-pool), 515 (partial trailing tile) and 257 (below one tile), on both entry points, plus the existing 2048 cases for batch 8192, context 1,000,000, zero row stride and a caller-provided output buffer
Registration	register_amd_ci(suite="stage-b-test-1-gpu-small-amd-mi35x") added beside the file's existing register_cuda_ci. The kernel is shared and the widths come from the model config, so both platforms want the coverage
pre-commit	validate registered test CI registries reports the same 582 pre-existing taxonomy violations as pristine main, i.e. none added

---

CI States

Latest PR Test (Base): ❌ Run #34937535760
Latest PR Test (Extra): ❌ Run #34937535557
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.