<!-- .github/pull_request_template.md -->

## 📌 Description

<!-- What does this PR do? Briefly describe the changes and why they’re
needed. -->

> Part of the FilteredTopK implementation refers to or is adapted from
@Linda-Stadter's work in #2759

### Deterministic Mode for Top-K Kernels

#### FilteredTopK Kernel

FilteredTopKKernel implements deterministic mode as follows:

1. Build a coarse histogram.

- Build a coarse histogram on the top 8 bits to locate the coarse
threshold bin that contains the k-th largest element.
- Same as non-deterministic mode, elements with bin > threshold_bin are
appended to s_indices via **atomicAdd** (see
`collect_gt_and_nondet_eq_threshold`); their final order is determined
by the post-sort kernel.

2. Refine with 8-bit radix passes.

- Run multiple 8-bit refine passes to find the exact pivot.
- Deterministic == pivot selection is performed by
`collect_det_eq_pivot`, which writes the selected tie elements into
`s_indices` in deterministic **thread-strided order**.

> **Thread-strided order** means, for example, if `BLOCK_THREADS = 4`,
then the logical scan order is:
>
> - thread 0: `0, 4, 8, ...`
> - thread 1: `1, 5, 9, ...`
> - thread 2: `2, 6, 10, ...`
> - thread 3: `3, 7, 11, ...`
>
> If the `== pivot` positions are:
> - thread 0: `0, 8`
> - thread 1: `5`
> - thread 2: none
> - thread 3: `3, 7`
>
>  then the deterministic collection order is: [0, 8, 5, 3, 7].
> That is, we order elements first by thread ID, and then by each
thread's strided traversal order.

3. Post-sort kernels.

- After FilteredTopKKernel finishes, `SortTopKByIndexKernel` is applied
to produce index-ascending output and make the final ordering
deterministic (we use atomicAdd to collect > pivot at stage 1).
- If the Python API is called with sorted=True,
`StableSortTopKByValueKernel` is applied afterward to produce
value-descending output.

#### RadixTopK Kernel


1. RadixSelectFindPivot

- Finds `ordered_pivot`, which Stage 2 uses to determine whether an
element is >= `ordered_pivot`.
- Computes `cta_local_eq_count` and `cta_local_gt_count`, which Stage 2
uses to **determine** how many elements the current CTA may emit and
where each emitted element should be placed.

2. collect_indices (`RadixCollectIndicesDeterministic`)

RadixCollectIndicesDeterministic: after the pivot is known, assigns each
CTA a fixed output range, then writes all > pivot elements followed by
the required == pivot elements in a deterministic order.

Order definition:

- Emit > pivot elements first, then == pivot elements.
- For each category, earlier CTAs write to earlier output positions.
- Within each CTA, emit elements in thread-strided order.

### Benchmarks

machine: NVIDIA A100-PCIE-40GB

command: (fp32/fp16/bf16)
```bash
python -u benchmarks/bench_topk.py \
  --op all \
  --dtype fp32 \
  --deterministic \
  --compare-torch-deterministic \
  --input-pattern random
```

raw results:


[output.txt](https://github.com/user-attachments/files/26337712/output-group2-current-v1_5-20260330-101411.txt)
**Summary**


| dtype | geomean det slowdown vs non-det | geomean speedup vs torch.det
|
  | --- | ---: | ---: |
  | fp32 | 1.0992x | 1.7660x |
  | fp16 | 1.0777x | 1.3381x |
  | bf16 | 1.0745x | 1.3055x |


NOTE: FlashInfer deterministic **underperforms** PyTorch mainly on
short-sequence workloads. Importantly, this is not unique to the
deterministic path: FlashInfer non-deterministic top-k is also slower
than PyTorch in the same short-sequence regime. This suggests the gap is
primarily a short-sequence top-k issue rather than a
deterministic-specific regression. Optimizing short-sequence top-k, for
both non-deterministic and deterministic modes, is better treated as
future work.



## 🔍 Related Issues

close: #2584


## 🚀 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

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] 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](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

```
unittest I ran:
test_topk.py
test_sampling.py
test_logits_processor.py
```
## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Deterministic mode for top‑k and fused transforms (stable, repeatable
tie ordering) with API flag to enable deterministic outputs and stable
sorting behavior.

* **Benchmarks**
* Expanded benchmarking to compare deterministic vs nondeterministic
runs, pre-generated input patterns, DSA workload cases, and richer CLI
output.

* **Tests**
* Large suite of determinism and correctness tests (ties, multi‑CTA,
streams, sorted behavior, cache transitions).

* **Bug Fixes**
  * Improved runtime-error labeling and benchmark cache handling.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Yinzuo Jiang <jiangyinzuo@foxmail.com>
Signed-off-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com>
Co-authored-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com>