🤖 AI

Qrita: High-performance Top-k and Top-p using Pivot-based Truncation and Selection

The paper introduces Qrita, a high-performance, deterministic Top-k and Top-p sampling algorithm for large vocabularies that utilizes pivot-based truncation and selection to achieve up to 1.4x throughput improvement and 50% memory reduction compared to existing GPU kernels, leading to its adoption as the default sampler in vLLM.

Original authors: Jongseok Park, Sunga Kim, Alvin Cheung, Ion Stoica

Published 2026-05-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jongseok Park, Sunga Kim, Alvin Cheung, Ion Stoica

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a chef running a massive kitchen (a Large Language Model) with a pantry containing 100,000 different ingredients (the vocabulary). Every time you need to cook a dish (generate a word), you have to pick the best few ingredients from that giant pantry to ensure the meal tastes good but isn't boring.

In the world of AI, this picking process is called Top-k (picking the top k best ingredients) and Top-p (picking the smallest group of ingredients that adds up to a certain "flavor score").

The Problem: The Slow, Cluttered Kitchen

Currently, most kitchens handle this by dumping all 100,000 ingredients onto a giant counter, sorting them by "taste score" from best to worst, and then grabbing the top ones.

The Issue: Sorting 100,000 items is slow and messy. It takes up a huge amount of counter space (memory) and the chef's hands get tired (computation overhead).
The Alternative: Some chefs try to just grab a random handful of ingredients that might be good. But this is risky; sometimes you miss the best ingredients, and sometimes the dish turns out different every time you cook it (non-deterministic), which is bad if you need to replicate a recipe exactly.

The Solution: Qrita (The Smart Sous-Chef)

The paper introduces Qrita, a new, high-speed method for picking ingredients. Think of Qrita as a super-smart sous-chef who uses two clever tricks to skip the messy sorting entirely.

Trick 1: The "Gaussian Sigma-Truncation" (The Noise Filter)

Imagine that in your 100,000-ingredient pantry, 99,000 of them are just "noise" (like salt, sugar, and flour that are all roughly the same bland taste). Only a few hundred are the "star ingredients" (like truffles or saffron).

Instead of looking at every single jar, Qrita does a quick sniff test. It calculates the average taste and the "spiciness" (standard deviation) of the pantry. It then draws a line: "Anything below this line is just noise; we don't need to look at it."

The Result: It instantly throws away 99% of the pantry, leaving only the 200 or so interesting jars. This happens in a single, lightning-fast pass.

Trick 2: The "Quaternary Pivot Search" (The Four-Way Split)

Now, the chef has a small pile of 200 interesting ingredients. They still need to find the exact top 50.

Old Way: Check one by one (too slow) or split the pile in half (binary search).
Qrita's Way: Instead of splitting the pile in half, Qrita splits it into four sections at once. It asks: "Is the best ingredient in the first quarter, second, third, or fourth?"
The Bonus: It also has a special rule for duplicate ingredients. If three jars of "Saffron" have the exact same score, Qrita knows exactly how many to keep to ensure the recipe is identical every time (deterministic). This prevents the chef from getting stuck in an endless loop trying to decide between identical jars.

Why This Matters (The Results)

The authors built Qrita using a specialized tool called Triton (a language for programming graphics cards/GPUs) and tested it against the current best methods used by major AI engines like vLLM and SGLang.

Speed: Qrita is up to 1.4 times faster in real-world serving scenarios and up to 2 times faster in raw speed tests.
Memory: It uses half the memory because it doesn't need to store the whole sorted list of 100,000 items.
Accuracy: Unlike some fast methods that guess, Qrita gives you the exact same result as the slow, perfect sorting method. It doesn't change the output; it just finds it much faster.

The Bottom Line

Qrita is like upgrading from a chef who manually sorts every single spice in the world to a smart assistant who instantly ignores the boring stuff, splits the interesting stuff into four piles at once, and handles duplicates perfectly. It makes AI generation faster and more efficient without changing the quality of the answers.

Note: The paper mentions that Qrita is now the default method for the GPU path in vLLM, a popular tool for running AI models, and the code is available for others to use.

Technical Summary: Qrita - High-performance Top-k and Top-p using Pivot-based Truncation and Selection

Problem Statement

Efficient implementation of Top-k and Top-p sampling algorithms for Large Language Models (LLMs) with large vocabularies (often exceeding 100,000 tokens) remains a significant challenge on GPUs. Existing approaches face a trade-off between computational efficiency and output correctness:

Sorting-based methods: Used by engines like vLLM, these provide exact, deterministic results but incur high computational and memory overhead due to full-vocabulary sorting and strided memory accesses.
Stochastic methods: Libraries like FlashInfer and SGLang use rejection sampling to avoid sorting, offering high throughput. However, they return only a single sampled token rather than the full candidate set, making them incompatible with tasks requiring deterministic candidate sets, such as speculative decoding verification and reinforcement learning.
Statistical cutoff methods: Approaches like Top-nσ and Min-p eliminate sorting by using statistical thresholds (e.g., mean and standard deviation). While efficient, they alter the sampling distribution and output, leading to debates regarding generation quality and coherence.

There is a critical gap for an algorithm that delivers memory efficiency and strict deterministic behavior (returning the full candidate set) without sacrificing the exact semantics of Top-k and Top-p.

Methodology

The authors propose Qrita, a high-performance Top-k and Top-p algorithm for GPUs based on pivot-based truncation and selection. Qrita extends the pivot-based strategy of RTop-k (originally designed for Graph Neural Networks) to handle large LLM vocabularies while ensuring determinism. The method relies on two key techniques:

1. Gaussian-based $\sigma$ -Truncation

Recognizing that LLM logits typically follow a quasi-Gaussian distribution where high-probability tokens form outliers against a "noisy" background, Qrita employs a pre-filtering step.

Mechanism: The algorithm approximates the logit distribution as $Z \sim \mathcal{N}(\mu, \sigma^2)$ . It computes the mean ( $\mu$ ) and standard deviation ( $\sigma$ ) from a sample block and calculates a truncation threshold $t = \mu + \delta \cdot \sigma$ .
Lookup Tables: The coefficient $\delta$ is determined via precomputed lookup tables derived from the normal quantile function, tailored to the target $k$ or $p$ values.
Outcome: This step truncates the massive "noise" region of the vocabulary in a single pass. If the truncated set size $|Z_t|$ exceeds the target $k$ (a "truncation hit"), the search proceeds only on this smaller subset. If not, it falls back to the full logit set. This significantly reduces the search space for the subsequent pivot search.

2. Quaternary Pivot Search with Duplication Handling

To find the exact pivot value $\tau$ (where exactly $k$ elements are $\ge \tau$ ) efficiently and deterministically:

Quaternary Search: Instead of a binary search, Qrita evaluates three pivots (at 25th, 50th, and 75th percentiles of the current search range) in each iteration. This reduces the number of search iterations by half compared to binary search.
Duplication Handling: Standard pivot searches can fail or become non-deterministic when duplicate logit values exist (common in low-precision or long-tailed distributions). Qrita tracks the minimum logit value greater than the pivot ( $z_{mn}$ ) and its occurrence count ( $N_{zmn}$ ).
Termination Condition: The search terminates when a pivot $\tau_n$ is found such that the count of elements strictly greater than $\tau_n$ is less than $k$ , but the count of elements greater than or equal to the minimum duplicate value $z_{mn}$ is greater than or equal to $k$ . This allows the algorithm to precisely determine how many duplicate instances of the boundary value must be retained to satisfy the Top-k constraint, ensuring reproducible, deterministic outputs.

Implementation Details

Framework: Implemented using Triton, a GPU programming language.
Autotuning: Employs a "Tiered Triton Block Size Autotune" strategy. Since the vector size varies drastically between the full logit set and the truncated set, the compiler is allowed to select optimal block sizes (e.g., larger blocks for full sets, smaller for truncated sets) to maximize hardware utilization.
Pipeline: The algorithm processes batches in parallel, applying Gaussian truncation, followed by the quaternary search, and finally a duplication handling mask generation to produce the final output.

Key Contributions

Deterministic High-Performance Sampling: Qrita provides the full Top-k/Top-p candidate set (unlike FlashInfer/SGLang) while maintaining the exact semantics of sorting-based algorithms, making it suitable for speculative decoding and RL.
Search Space Reduction: By leveraging the statistical properties of logits, the Gaussian $\sigma$ -truncation drastically reduces the search space, avoiding full-vocabulary sorting.
Robust Pivot Search: The introduction of quaternary search combined with explicit duplication handling guarantees deterministic behavior and reduces iteration counts, addressing the inefficiencies of standard pivot-based methods on large vocabularies.
Integration: Qrita is implemented as the default Top-k and Top-p sampler for the GPU execution path of vLLM.

Results

The authors evaluated Qrita against baselines including naive PyTorch sorting, vLLM's default sampler, FlashInfer, and SGLang on NVIDIA H100, RTX4090, and AMD MI355X GPUs using various models (Llama, Qwen, GPT-OSS, Gemma).

Throughput: Qrita achieves up to 1.4 $\times$ end-to-end serving throughput improvement over the vLLM baseline on H100, with some configurations showing up to 2 $\times$ improvement on RTX4090. It outperforms sorting-based baselines significantly while matching the output of sorting-based algorithms.
Memory Usage: Qrita reduces memory usage by approximately 50% compared to sorting-based methods, primarily by avoiding the allocation of large intermediate buffers required for sorting.
Ablation Studies:
- Gaussian Truncation: Provides the largest speedup, reducing latency by ~74% compared to searching the full logit set.
- Quaternary Search: While slightly slower than binary search in "truncation hit" cases due to reduced search space, it provides a significant latency reduction (~2.3ms) in fallback cases where the full vocabulary must be searched.
- Duplication Handling: Adds minimal overhead (~3.5% latency increase) but is critical for correctness and determinism.
- Autotuning: Reduces execution latency by nearly 30% compared to using a single fixed block size.

Significance and Claims

The paper claims that Qrita bridges the gap between memory efficiency and strict deterministic behavior for large-vocabulary decoding. It asserts that by avoiding full-vocabulary sorting and leveraging pivot-based selection with statistical pre-filtering, Qrita offers a high-performance solution for LLM decoding that is compatible with advanced tasks requiring full candidate sets (e.g., speculative decoding). The authors position Qrita not as a replacement for the statistical semantics of Top-k/Top-p, but as an efficient implementation that preserves them exactly, unlike heuristic statistical cutoff methods or stochastic single-token samplers. The integration into vLLM as the default sampler underscores its practical utility in production environments.