fix: topK uint32 overflow

[Linda-Stadter]

📌 Description

This PR is based on top of #2661. Merge afterwards.

Fixes uint32 overflow in top_k when batch_size * vocab_size > 2^32

The row offset row_idx * stride was computed in uint32 arithmetic, silently
wrapping to zero for large inputs. Cast to size_t before the multiplication.

Add a regression test with batch=32769, vocab=131072 (fp16) that crosses the
overflow boundary.

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: fd690d55-f195-46b2-9e2a-114315775aa3

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request introduces deterministic mode for FlashInfer's top-k operations, including basic top-k selection and fused transforms (page table and ragged). The implementation includes new CUDA kernels for deterministic index collection and post-processing stable sorts to ensure bitwise-reproducible results across different runs and environments. The Python API has been updated to expose a deterministic flag, and the benchmarking suite is expanded to include DeepSeek DSA-like workloads and various input patterns to stress-test tie-handling logic. Extensive tests have been added to verify repeatability and correctness. I have no feedback to provide as there were no review comments to assess.

[aleozlx]

cherry picked into #2661
closing for now