[AMD] [GLM-5.3-Flash Day 0] Accept k-pool top-k widths in the DSA page-table transform

[Raiden-Makoto]

Summary

Replacement for #37573, which GitHub closed when its deleted support base was merged into main. The original author and review history remain credited there.

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_TOPK	Kernels	When
true	one — fast_kpool_topk_transform_fused does the top-k and the page-table gather, so this PR's entry points are never reached	the default
false	two — the same kernel returns raw indices, then transform_index_page_table_{decode,prefill}_fast gathers	--dsa-topk-backend torch; PD disaggregation with MTP index share, where should_use_dsa_fused_topk() returns false whatever the env says; or set by hand to bisect

Both of those entry points 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 reject it.

Scope

transform_index.py only. No environment variable, no architecture gate.

Case	Effect
width 2048	decode keeps the existing single-program kernel, prefill is untouched — every current caller keeps its kernel and its performance
other widths	the new tiled kernel transform_index_page_table_decode_tiled_kernel
platform	AMD and NV. CUDA raises the same AssertionError under the same configuration

Test plan

Combined main-based integration (RM/glm53-day0-main-integration tested at 9d0d62b744; current integration head 4e2f84ff3c; TP8/TP4 validation predates the above-cap guard, whose below-cap path is unchanged, pinned AITER d9e5ef7ce0, fresh caches): TP8 GSM8K 97.50% in 87.0s with 0.00% errors; TP4 GSM8K 96.97% in 320.2s with 0.00% errors. Decode CUDA graphs were enabled in both runs.

Recovery base: main @ 30e7a3072d.
Recovery validation on MI355X: scoped pre-commit passed; 11 transform-index tests passed.

Docker: 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.

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	all hooks pass on both files, including validate registered test CI registries
Accuracy	not measured — GLM-5.3-Flash still cannot be served on ROCm from this branch; GSM8K lands once the remaining enablement is in

---

CI States

Latest PR Test (Base): ⏳ Run #34768735944
Latest PR Test (Extra): ❌ Run #34768735764
Latest PR Test (AMD ROCm 10): ❌ Run #34768735873

[Raiden-Makoto]

Closing this replacement because @Jacob0226 will reopen and maintain the original PR, #37573. Please continue review and CI there.

[Jacob0226]

Thanks @Raiden-Makoto. Reopening as #39340 rather than #37573: #37573's base branch xinyuan/glm-5.3-flash-support has been deleted, so GitHub will not let it reopen. #39340 is off main and its changed files are byte-identical to this PR's head.