[sgl-kernel][Deepseek V3.2] Add row_starts to topk kernel

[hlu1]

Motivation

Part1 of the fix for bug in #11629

Modifications

In the topk kernel for prefill, the q and k inputs are both ragged. We need to pass the correct start indices of k for each q token to the kernel.

In fast_topk_transform_interface, I changed the is_decode criterion from prefill_bs == B to !row_starts_opt.has_value() and prefill_bs == B. Using prefill_bs == B can be a bit ambiguous for a single token.

Accuracy Tests

8xB200, fp8 checkpoint, bf16 kvcache, radix-cache on:
fp8 kvcache has similar score.

gsm8k

20 shots

Accuracy: 0.953
Invalid: 0.000
Latency: 23.981 s
Output throughput: 5388.611 token/s

30 shots

Accuracy: 0.928
Invalid: 0.000
Latency: 31.716 s
Output throughput: 3907.670 token/s

gpaq

Repeat: 8, mean: 0.790

['0.798', '0.778', '0.783', '0.803', '0.768', '0.813', '0.788', '0.788']

With mtp on

gsm 8k with 20 shots: 0.947
gsm 8k with 30 shots: 0.939

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.

[gemini-code-assist]

Summary of Changes

Hello @hlu1, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical issue in the 'topk' kernel's handling of ragged inputs, which is essential for the Deepseek V3.2 model's prefill stage. By integrating a 'row_starts' parameter, the kernel can now accurately determine the starting indices for each row of a ragged tensor, ensuring correct top-k computation. This enhancement improves the robustness and correctness of the model's operations, directly contributing to resolving a reported bug.

Highlights

• Ragged Input Support: Introduced a 'row_starts' parameter to the 'fast_topk' CUDA kernel and its associated transform functions to correctly handle ragged inputs, particularly for prefill operations in the Deepseek V3.2 model.
• API Extension: Modified the C++ and Python interfaces for 'fast_topk', 'fast_topk_transform_fused', and 'fast_topk_transform_ragged_fused' to accept the new 'row_starts' tensor, making it optional for some and required for others.
• Bug Fix (Part 1): This change is identified as the first part of a fix for a known bug related to incorrect indexing in the top-k kernel when dealing with ragged inputs.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request correctly adds the row_starts parameter to the top-k kernels to support ragged tensors, which is a good improvement. The changes are propagated through the C++ kernels, PyTorch op definitions, and Python wrappers. My review has identified a few areas for improvement: there are a couple of missing input validation checks in the C++ code, a minor typo in a Python assertion message, and most importantly, the tests for the new functionality are incomplete as they don't cover cases with non-zero row_starts. I've provided suggestions to address these points.

[Fridge003]

@hlu1 Can you try building the kernels with make build under sgl-kernel folder, and see whether this PR along with #12583 can fix this bug?

[hlu1]

@hlu1 Can you try building the kernels with make build under sgl-kernel folder, and see whether this PR along with #12583 can fix this bug?

That's exactly how I tested it. The 20shots gsm8k results are now matching vllm results in https://docs.vllm.ai/projects/recipes/en/latest/DeepSeek/DeepSeek-V3_2-Exp.html#additional-resources

[Fridge003]

@hlu1 We can add the 20-shot test to CI test after your PRs are merged, for testing long context accuracy