Further optimize top-k kernel for DSv3.2

[bkryu]

FlashInfer currently has a radix topk that is somewhat well optimized.

DSv3.2's lightning indexer requires a top-k selection kernel tuned for k=2048. Would like to have a further optimized top-k kernel that can be integrated to the frameworks.

[bkryu]

SGLang has three top-k variants:

• (unfused) top-k -- equivalent to our top-k
• fast_topk_transform_fused -- equivalent to our top_k_page_table_transform
• fast_topk_transform_ragged_fused -- equivalent to our top_k_ragged_transform

Need to look into what to support.

[huangzhilin-hzl]

SGLang has three top-k variants:

• (unfused) top-k -- equivalent to our top-k
• fast_topk_transform_fused -- equivalent to our top_k_page_table_transform
• fast_topk_transform_ragged_fused -- equivalent to our top_k_ragged_transform

Need to look into what to support.

Hi @bkryu
I've noticed a few key differences between the two implementations:

• sgl_kernel.fast_topk_v2 supports variable-length inputs (lengths), while FlashInfer does not.
• sgl_kernel.fast_topk_transform_fused supports cases where score.shape[1] > src_page_table.shape[1]. In FlashInfer, this scenario triggers an illegal memory access. Could this be a bug? Ref: [top_k_page_table_transform bug]Illegal memory access in RadixTopKKernel_Unified #2486
• The core implementation of radix select is conceptually similar in both. However, FlashInfer's implementation is more comprehensive, featuring vectorized reads, support for multiple score dtypes, and multi-CTA collaboration. This is why FlashInfer's implementation is significantly faster than SGLang's in very long-context scenarios.

Based on this, I was wondering if there's still room for further optimization. I have a couple of ideas:

• Using lower precision for top-k computation. For example, we could convert the score tensor to FP16 right before the top-k calculation. The rationale is that building the histogram requires repeatedly scanning each row of the score data. A lower-precision format would allow LDG instructions to load more data at once and could also reduce the number of execution rounds for radix select, thereby lowering the overall kernel latency. Of course, this would need to be tested for its impact on numerical precision and end-to-end model accuracy.
• Other optimization techniques. For instance, could we use persistent blocks when num_rows is much larger than the number of SMs? Or are there further optimizations possible for atomicAdd?

[yzh119]

@huangzhilin-hzl

supports variable-length inputs (lengths), while FlashInfer does not.

It should be easy to add support, we have enabled variable length in top_k_page_table_transform and top_k_ragged_transform, we just need to do the same thing for flashinfer.top_k.

Could this be a bug? Ref: #2486

Thanks for the bugfix, I'll take a look.

Regarding optimizations, @limin2021 is working on this.

For instance, could we use persistent blocks when num_rows is much larger than the number of SMs? Or are there further optimizations possible for atomicAdd?

This is where multi-cta helps, our multi-cta algorithm didn't use the filtering algorithm which we can further optimize.

[huangzhilin-hzl]

@yzh119

It should be easy to add support, we have enabled variable length in top_k_page_table_transform and top_k_ragged_transform, we just need to do the same thing for flashinfer.top_k.

Is anyone already working on a solution for this variable-length issue? If not, I'd be willing to contribute a pr.