flashinfer.sdsa.thresholds_from_sample

flashinfer.sdsa.thresholds_from_sample(sample_scores: Tensor, sample_top: Tensor, *, thresholds: Tensor, counters: Tensor) → None

Set each query row’s threshold for indexer_threshold_select() from a sample of its scores, and zero its counters. It runs as one launch without host synchronization, so it can be captured in a CUDA graph.

Parameters:

• sample_scores (torch.Tensor) – float32 [rows, samples]: each row’s sampled scores.

• sample_top (torch.Tensor) – int32 [rows, rank]: each row’s Top positions among its samples.

• thresholds (torch.Tensor) – float32 [rows], output: the least of the row’s sampled Top scores. A -1 position or a NaN score counts as -inf, since dummy and capture passes can feed NaN from uninitialized memory.

• counters (torch.Tensor) – int32 [rows, 2], output: zeroed.