[Feature] Add optional TileLang Top-K selector

[qiyueliuhuo23]

Summary

This PR adds an optional TileLang radix Top-K implementation for the GLM-5.2 DSA indexer path.

PyTorch topk remains the default behavior. The new implementation is enabled only when:

  export USE_TILELANG_TOPK=1

Background

The GLM-5.2 DSA indexer first produces an FP32 logits matrix and then selects the Top-K KV positions using torch.topk.

For packed causal sequences, each query only has a valid [start, end) KV range. However, the original PyTorch path still scans the complete logits row, including invalid positions that have already been masked to -inf.

The TileLang implementation uses these valid ranges directly to avoid unnecessary processing outside the causal interval.

Main Changes

1. Add TileLang radix Top-K

New file:

  xtuner/v1/ops/sparse_mla/tilelang_topk.py

The implementation:

• Uses FP32 radix keys for Top-K selection.
• Accepts per-row starts and ends ranges.
• Uses an 8192-entry candidate buffer.
• Falls back to dense rescan when the radix boundary bucket exceeds the buffer capacity.
• Avoids out-of-bounds writes for concentrated score distributions.
• Reorders selected indices by score to match torch.topk(sorted=True).
• Returns int32 indices and uses -1 for invalid causal positions.

2. Add an independent runtime switch

Default behavior remains PyTorch Top-K:

  export USE_TILELANG_TOPK=0

Enable TileLang Top-K explicitly:

  export USE_TILELANG_TOPK=1

The switch is exposed through:

  DSAMLAConfig.use_tilelang_topk

It only changes the Top-K selection stage. It does not change:

• The indexer GEMM;
• Sparse MLA forward or backward backends;
• The default training path;
• The PyTorch reference implementation.

3. Add correctness tests

New test file:

  tests/ops/test_tilelang_topk.py

The tests cover:

• Normal score distributions;
• Concentrated score distributions;
• Packed causal ranges;
• Rows with fewer valid values than topk;
• Bitwise equality between TileLang and PyTorch Top-K outputs;
• The integrated indexer switch path.

The tests use unique FP32 values so that the comparison is not affected by PyTorch's unspecified ordering for tied scores.

Standalone Performance

The measurement includes the complete selector path:

• TileLang radix selection;
• Padding handling;
• Selected-score gather;
• Output sorting;
• Index reordering.

Pack length	PyTorch median	TileLang median	Speedup
16K	4.250432ms	2.660752 ms	1.5975x
32K	14.518560ms	5.570000ms	2.6066x
64K	56.546545ms	12.461776ms	4.5376x
128K	214.148964ms	30.375248ms	7.0501x

The main benefit comes from skipping invalid causal ranges. For long single-segment inputs or highly concentrated distributions, the dense rescan fallback can reduce or eliminate the benefit.

End-to-End Training Performance

Configuration:

• 8 × H200;
• EP4, SP1;
• 16K packed sequence;
• 30 warm-up steps + 100 measured steps;
• Baseline: PyTorch Top-K;
• Candidate: TileLang Top-K;
• All other settings kept unchanged.

Metric	Baseline	Candidate	Change
Step-time median	4.547550 s	4.537700 s	-0.2166%
Step-time P90	4.565840 s	4.559440 s	-0.1402%
Step-time P99	4.604562 s	4.620483 s	+0.3458%
TGS	28,509.145	28,557.332	+0.1690%
Peak allocated memory	104.582 GiB	104.885 GiB	+0.303 GiB

The official CI checker passed:

• loss/local_loss maximum relative difference: 0.4307%;
• loss/reduced_llm_loss maximum relative difference: 0.2531%;
• grad_norm maximum relative difference: 40.3339%;
• TGS P80 degradation: 5.7547%, below the 6% threshold;
• Memory, learning rate, and token-count checks passed.

Expected End-to-End Benefit

The current Top-K path accounts for approximately 0.38% of the total step time.

Using the standalone speedup range of 1.47x–1.75x:

Expected step-time improvement
= 0.38% × (1 - 1 / speedup)
≈ 0.121%–0.163%

The measured average step-time improvement was 0.1687%, and TGS improved by 0.1690%, which is consistent with the theoretical estimate.

Compatibility

• USE_TILELANG_TOPK=0 preserves the existing PyTorch behavior.
• No additional external dependency is required.
• TileLang Top-K is imported only when explicitly enabled.
• The torch and cudnn_dsa indexer backends do not automatically switch to TileLang Top-K.
• The implementation is intended for the TileLang indexer path.

Validation

  export USE_TILELANG_TOPK=1
  pytest -q tests/ops/test_tilelang_topk.py

To use the default PyTorch implementation:

  export USE_TILELANG_TOPK=0