FlashTrie: A GPU-Accelerated Constrained Beam Search for Generative Retrieval

# FlashTrie: A GPU-Accelerated Constrained Beam Search for Generative Retrieval
Source: [https://arxiv.org/html/2607.10044](https://arxiv.org/html/2607.10044)
Dakshitha Anandakumar1Anurag Mukkara2Wenxiang Hu1Jiusheng Chen1 M Akash Kumar1Ting Ye1Qiang Lou1Jian Jiao1 1Microsoft, Redmond, WA, USA2Nvidia, Santa Clara, CA, USA \{danandakumar, tiy\}@microsoft\.com

###### Abstract

Constrained decoding is essential in generative retrieval, where document identifiers generated directly from a query must exactly match a predefined library of valid IDs\. At scale, decoding is often constrained using a trie with beam search but most implementations run on CPU\. Limited parallelism then makes trie traversal and candidate validation a serving bottleneck as beam width grows\.

We presentFlashTrie, which addresses this limitation by optimizing constrained beam search on GPUs\. It introduces an integer\-aware succinct trie layout that uses bit compression to reduce memory footprint while keeping the full index in GPU high\-bandwidth memory reducing memory stalls, and a cooperative CUDA kernel that performs beam expansion, validation, and pruning entirely on\-device without per\-step host orchestration\. It further replaces CPU\-style irregular lookup and heap maintenance with GPU\-aware parallel primitives, improving warp utilization and reducing divergence\.

Together, these designs significantly reduce decoding latency and increase throughput while preserving retrieval quality\. On a library of 800M keywords with beam widths up to10001000, FlashTrie reduces trie\-search latency to under33ms, achieving up to24×24\\timesspeedup over a highly optimized multi\-threaded CPU baseline\. These improvements enable FlashTrie to scale beam sizes by up to5×5\\timesin latency\-critical applications such as sponsored search\. In a large\-scale online A/B experiment on a commercial search engine, it delivers a statistically significant\+0\.71%\+0\.71\\%revenue lift, enabling real\-time constrained decoding at a scale previously feasible only offline\. The FlashTrie code will be publicly released after the review process\.

FlashTrie: A GPU\-Accelerated Constrained Beam Search for Generative Retrieval

Dakshitha Anandakumar1Anurag Mukkara2Wenxiang Hu1Jiusheng Chen1M Akash Kumar1Ting Ye1Qiang Lou1Jian Jiao11Microsoft, Redmond, WA, USA2Nvidia, Santa Clara, CA, USA\{danandakumar, tiy\}@microsoft\.com

## 1Introduction

Generative retrievalTayet al\.\([2022](https://arxiv.org/html/2607.10044#bib.bib1)\); Metzleret al\.\([2021](https://arxiv.org/html/2607.10044#bib.bib54)\)reframes document retrieval as a sequence\-to\-sequence task, mapping a query directly to an identifier \(docID\), replacing dual\-encoder indexingKarpukhinet al\.\([2020](https://arxiv.org/html/2607.10044#bib.bib36)\)\. This paradigm is attractive for large\-scale search and recommendation systems because it can, in principle, avoid expensive retrieval pipelines while enabling compact end\-to\-end modeling\. In practice, however, online serving is constrained by strict latency budgets, and the decoding procedure becomes the dominant bottleneck\. Autoregressive \(AR\)Sutskeveret al\.\([2014](https://arxiv.org/html/2607.10044#bib.bib55)\)decoders are sequential and costly, while non\-autoregressive \(NAR\)Guet al\.\([2018](https://arxiv.org/html/2607.10044#bib.bib12)\); Sun and Yang \([2020](https://arxiv.org/html/2607.10044#bib.bib37)\)decoders regain parallelism but their per\-position independence frequently produce invalid identifiers that must be filtered by constrained searchZiemset al\.\([2023](https://arxiv.org/html/2607.10044#bib.bib7)\); Pradeepet al\.\([2023](https://arxiv.org/html/2607.10044#bib.bib39)\)\.

#### Prior Work and Challenges\.

Several strategies restrict decoding to valid outputs, including logit maskingTayet al\.\([2022](https://arxiv.org/html/2607.10044#bib.bib1)\), finite\-state compilationWillard and Louf \([2023](https://arxiv.org/html/2607.10044#bib.bib11)\), and predicate\-logic frameworksLuet al\.\([2021](https://arxiv.org/html/2607.10044#bib.bib10)\); Andersonet al\.\([2017](https://arxiv.org/html/2607.10044#bib.bib40)\)\. As the identifier library scales to millions or billions of entries, a common approach is trie\-constrained beam searchHokamp and Liu \([2017](https://arxiv.org/html/2607.10044#bib.bib25)\), used across flat\-token schemes \(DSITayet al\.\([2022](https://arxiv.org/html/2607.10044#bib.bib1)\), GENRECaoet al\.\([2021](https://arxiv.org/html/2607.10044#bib.bib2)\)\), and structured semantic\-identifiers \(NCIWanget al\.\([2022](https://arxiv.org/html/2607.10044#bib.bib6)\), SEALBevilacquaet al\.\([2022](https://arxiv.org/html/2607.10044#bib.bib5)\), and recent workPenhaet al\.\([2025](https://arxiv.org/html/2607.10044#bib.bib43)\)\)\.

However, pointer\-based trie representationsMorrison \([1968](https://arxiv.org/html/2607.10044#bib.bib17)\); Aoe \([1989](https://arxiv.org/html/2607.10044#bib.bib16)\)suffer from irregular memory access and poor hardware utilization, leading production systems to rely on optimized CPU implementationsCaoet al\.\([2021](https://arxiv.org/html/2607.10044#bib.bib2)\)\. Succinct designs such as MARISAYata \([2011](https://arxiv.org/html/2607.10044#bib.bib14)\), built on LOUDSJacobson \([1989](https://arxiv.org/html/2607.10044#bib.bib15)\)and minimal acyclic FSAsDaciuket al\.\([2000](https://arxiv.org/html/2607.10044#bib.bib18)\), improve space efficiency but support only character\-level keys, limiting their applicability to generative retrieval with large token vocabularies and batched, beam\-aware traversal\. These structures also map poorly to GPUs, causing warp divergence and uncoalesced accessMerrillet al\.\([2012](https://arxiv.org/html/2607.10044#bib.bib19)\)\.

While modern inference runtimesWanget al\.\([2021](https://arxiv.org/html/2607.10044#bib.bib21)\); Daoet al\.\([2022](https://arxiv.org/html/2607.10044#bib.bib26)\); Kwonet al\.\([2023](https://arxiv.org/html/2607.10044#bib.bib22)\)optimize model execution on GPU, constraint enforcement is often implemented outside the core GPU decoding path, which can introduce additional orchestration overhead\. This motivates three key questions:*\(i\)*can succinct tries be adapted to token\-level vocabularies without sacrificing space efficiency;*\(ii\)*can constrained beam search be executed entirely on GPU; and*\(iii\)*can such systems meet production latency constraints without degrading retrieval quality?

#### Our Contribution\.

We presentFlashTrie, a GPU\-native constrained decoding framework for generative retrieval that jointly optimizes constraint representation and decoding computation\. Building on succinct tries, we redesign MARISA for integer\-token vocabularies using a bit\-compressed layout that reduces index size and keeps the constraint structure resident in GPU high\-bandwidth memory \(HBM\)\. We further introduce GPU\-friendly search and sort primitives, parallel trie child matching and heap\-free beam selection, to replace pointer\-heavy traversal and reduce divergence\. Constrained beam search is executed fully on\-device via a cooperative multi\-step CUDA kernel that performs expansion, validation, and pruning without repeated kernel relaunch overhead across decoding steps\. Two\-level parallelism, beam\-parallel and top\-KK\-parallel execution, exploits hierarchical GPU concurrency, saturating hundreds of streaming multiprocessors \(SMs\) and maintaining high occupancy across the device\. FlashTrie scales constrained decoding to billion\-scale constraint libraries on GPU, a regime previously feasible only in offline settings, while achieving up to 24×\\timesspeedup over optimized CPU baselines and reducing trie\-search latency to under 3 ms\. In production A/B testing on a large\-scale sponsored\-search system, these latency gains enable wider beam search within strict serving budgets, translating to a \+0\.71% revenue lift\.

## 2Design and Implementation

We consider constrained decoding with a trie𝒯\\mathcal\{T\}over vocabulary𝒱\\mathcal\{V\}\. FlashTrie applies to both AR and NAR decoding; we focus on NAR in experiments because larger branching factors make constraint checking more expensive\. At decoding stepttthe base model emits top\-KKcandidates\(xt,i,ℓt,i\)\(x\_\{t,i\},\\ell\_\{t,i\}\)withℓt,i=log⁡p​\(xt,i∣x<t\)\\ell\_\{t,i\}=\\log p\(x\_\{t,i\}\\mid x\_\{<t\}\); trie\-constrained beam search keeps the top\-BBvalid prefixes under length\-normalized log\-probabilityWuet al\.\([2016](https://arxiv.org/html/2607.10044#bib.bib13)\)\(Equation[2](https://arxiv.org/html/2607.10044#A5.E2), Appendix[E](https://arxiv.org/html/2607.10044#A5)\)\. For comparison, we use two CPU baselines built on the same LOUDS\+tail/link trie skeletonYata \([2011](https://arxiv.org/html/2607.10044#bib.bib14)\)\(Appendix[A](https://arxiv.org/html/2607.10044#A1)\)\.

### 2\.1CPU Baselines: MARISA\-Int and MARISA\-Opt

MARISA\-Intextends open\-source MARISAYata \([2011](https://arxiv.org/html/2607.10044#bib.bib14)\)from character\-level keys to 32\-bit integer tokens and adds a minimal beam\-search layer, while preserving the original build pipeline\. It runs single\-threaded on CPU\.MARISA\-Optis a stronger CPU baseline: adds trie traversal with prefix\-state caching, a multithreaded worker pool, two\-pointer merge between sorted proposals and child labels \(O​\(K\+Δ\)O\(K\{\+\}\\Delta\)per beam\), and a bounded min\-heap for top\-BBpruning \(O​\(C​log⁡B\)O\(C\\log B\)\)\. Design details and per\-operation complexity are in Appendices[B](https://arxiv.org/html/2607.10044#A2)and[C](https://arxiv.org/html/2607.10044#A3)\.

### 2\.2FlashTrie: GPU\-Resident Constrained Beam Search

![Refer to caption](https://arxiv.org/html/2607.10044v1/x1.png)Figure 1:Overview of FlashTrie pipeline\. At each decoding step, active beam states are mapped to thread blocks \(CTAs\) and processed in parallel across the threads of a CTA for trie\-based validation and beam expansion all executed on\-device\. \(Inset\) FlashTrie extends Narrow\-LOUDS representation to integer tokens stored in compact structured arrays\.FlashTrie redesigns both the trie data structure and the decoding algorithm for GPU execution\. The result is a GPU\-resident constraint index that preserves succinct\-trie space efficiency while reducing memory stalls and exposing hierarchical parallelism\.

#### Narrow LOUDS for integer tokens

Published MARISA targets 8\-bit alphabets\. FlashTrie extends the layout to 32\-bit integer tokens to enable large vocabularies \(over a million\) but narrows the per\-node slot by splitting both label and suffix index into a 2\-byte low part in the base array and a bit\-packed high part in a secondary array for nodes that need them\. This halves the dominant array, shrinks the constraint index, and keeps billion\-scale libraries in GPU HBM \(Figure[1](https://arxiv.org/html/2607.10044#S2.F1)\)\. It also improves locality and reduces pointer\-chasing stalls as fewer bits are fetched per access \(Section[4\.1](https://arxiv.org/html/2607.10044#S4.SS1)\)\.

#### Kernel Execution and Parallelism\.

FlashTrie executes constrained beam search in a single cooperative kernel, avoiding per\-step CPU\-GPU orchestration\. The GPU grid has several multiprocessors \(SM\) with each SM having one cooperative thread block \(CTA\)NVIDIA Corporation \([2024](https://arxiv.org/html/2607.10044#bib.bib3),[2020](https://arxiv.org/html/2607.10044#bib.bib4)\)\. Queries are batched across the SMs\. To enable parallel beam search, within each query, an active beam is assigned to one CTA\. A CTA has 512 threads that perform top\-KKexpansion and trie\-based validation via batched child lookup\. A grid\-synchronized step then merges and prunes beams using parallel merge sort to enforce beam width\.

At each decoding step, FlashTrie exposes two levels of parallelism\. First, a beam\-parallel expansion across allB​WBWactive beams\. Second, proposal\-parallel validation within each beam, where each thread tests one of the top\-KKlanguage\-model proposals through binary search over sorted trie child labels\. The fullTT\-step loop runs in one GPU launch \(Algorithm[1](https://arxiv.org/html/2607.10044#alg1)\), hence atB​W=K=1000BW=K=1000, this yields up to one million concurrent child\-lookups across the GPU grid\.

After expansion, the CPU heap is replaced by an append\-only candidate buffer and a parallel top\-BBselection kernel\. This removes lock contention and pointer chasing while keeping intermediate candidates in fast on\-chip memory \(Appendix[D](https://arxiv.org/html/2607.10044#A4)\)\.

Algorithm 1FlashTriebeam search\. One persistent GPU kernel executes allTTsteps on device, exposing beam\-parallel expansion and proposal\-parallel validation\. Kernel details: Appendix[D](https://arxiv.org/html/2607.10044#A4)\.0:trie

𝒯\\mathcal\{T\}, initial state

s0s\_\{0\}, beam width

B​WBW, length

TT, per\-step LM proposals

TopK​\(s,K\)\\textsc\{TopK\}\(s,K\)
1:

𝑐𝑢𝑟←\{s0\}\\mathit\{cur\}\\leftarrow\\\{s\_\{0\}\\\}
2:for

t=0t=0to

T−1T\{\-\}1do

3:

𝑛𝑒𝑥𝑡←∅\\mathit\{next\}\\leftarrow\\emptyset
4:foreach

s∈𝑐𝑢𝑟s\\in\\mathit\{cur\}in paralleldo

5:\{\(i\) beam\-parallel\}

6:

\(τ1:K,ρ1:K\)←TopK​\(s,K\)\(\\tau\_\{1\{:\}K\},\\rho\_\{1\{:\}K\}\)\\leftarrow\\textsc\{TopK\}\(s,K\)
7:for

k=1k=1to

KKin paralleldo

8:\{\(ii\) proposal\-parallel\}

9:

u←ChildLookup\(𝒯,s\.v,τk\)u\\leftarrow\\textsc\{ChildLookup\}\(\\mathcal\{T\},s\.v,\\tau\_\{k\}\)
10:\{

O​\(log⁡Δ\)O\(\\log\\Delta\)binary search\}

11:if

u≠⊥u\\\!\\neq\\\!\\botand

s\.σ\+ρk\>θsents\.\\sigma\{\+\}\\rho\_\{k\}\>\\theta\_\{\\mathrm\{sent\}\}then

12:append

\(s,u,ρk\)\(s,u,\\rho\_\{k\}\)to

𝑛𝑒𝑥𝑡\\mathit\{next\}
13:\{lock\-free\}

14:endif

15:endfor

16:endfor

17:barrier;

𝑐𝑢𝑟←Top\-​B​\(𝑛𝑒𝑥𝑡\)\\mathit\{cur\}\\leftarrow\\textsc\{Top\-\}B\(\\mathit\{next\}\)
18:\{on\-device\}

19:endfor

20:

21:returnBacktrace\(

𝑐𝑢𝑟\\mathit\{cur\}\)

## 3Experimental Setup

We evaluate on 13,000 retrieval requests from a production NAR model with a 2\.2M\-token vocabulary\. Each request generates top\-KKproposals \(∈100,…,1000\\in\{100,\\ldots,1000\}, beamwidth \(BW\)=K=K\) overT=8T=8decoding steps\. The constraint trie is built with 800M keyword sequences and is resident on a single A100 80GB GPU paired with an AMD EPYC 7V13 host\. For MARISA\-Opt, we use 8 CPU workers to saturate the 8\-core allocation; full hardware isolation details in Appendix[E](https://arxiv.org/html/2607.10044#A5)\.

We measure per\-request latency \(mean, p50/p90/p95/p99\), throughput under batched serving \(b∈1,4,8,16,32b\\in\{1,4,8,16,32\}queries\), on\-disk index size and wall\-clock build time, and Precision@100/200 scored by a transformer teacher \(Section[4\.5](https://arxiv.org/html/2607.10044#S4.SS5)\)\. Latency timing spans the full request path \(C\+\+ entry to result return, including GPU stream synchronization\)\. All timings average 10 passes over the full 13,000\-request set after one warm\-up discard\. Details in Appendix[E](https://arxiv.org/html/2607.10044#A5)\.

## 4Results

### 4\.1Trie Build Time and Index Size

We begin by evaluating how FlashTrie’s storage redesign affects two immediate construction outcomes: trie build time and on\-disk index size\.FlashTrie reduces both as the constraint set grows \(Figure[2](https://arxiv.org/html/2607.10044#S4.F2)\)\. Across 10M–800M sequences, the widening gap reflects a shift in the dominant cost: at small scales, fixed LOUDS and tail\-construction overheads dominate \(Appendix[A](https://arxiv.org/html/2607.10044#A1)\), whereas at larger scales the per\-edge link machinery and traversal bookkeeping become the bottleneck \(Appendix[C](https://arxiv.org/html/2607.10044#A3)\)\. MARISA\-Int and MARISA\-Opt remain close because they retain the same full\-width node storage layout, while FlashTrie reduces per\-node work by narrowing link offsets and splitting each 32\-bit label into a compact base field plus packed high bits \(Section[2\.2](https://arxiv.org/html/2607.10044#S2.SS2)\)\. As the corpus grows, those savings compound\. Even at such a large constraint index as 800M, trie size is only around 3 GB in FlashTrie, allowing the trie to fit entirely in GPU HBM and enabling deployment on smaller GPUs like T4 or A100 MIG\. Further, FlashTrie achieves 3\.4×\\timesfaster build time than MARISA\-Opt and uses 22% less index space\. Taken together, these construction\-time gains establish the practical value of the storage redesign, which is one of the key enablers of the latency results that follow\.

![Refer to caption](https://arxiv.org/html/2607.10044v1/x2.png)Figure 2:Trie construction scaling on 10M–800M keyword sequences, \(Left\) Build time, \(Right\) Constraint index size on disk; at runtime, the full trie is resident in GPU memory\. At 10M keys, FlashTrie builds in 0\.15 min and uses 0\.061 GB, compared with 0\.48 min / 0\.060 GB for MARISA\-Opt and 0\.58 min / 0\.075 GB for MARISA\-Int\. At 800M keys, FlashTrie reaches 13\.4 min and 3\.1 GB, versus 45\.3 min / 4 GB for MARISA\-Opt and 51\.1 min / 4 GB for MARISA\-Int\.
### 4\.2Latency Results

Next, we evaluate how GPU occupancy\-focused kernel optimizations, together with full on\-device residency of the constraint index, impact trie\-search latency as beam width increases\. On the GPU, each request runs in a single persistent cooperative kernel that executes every decoding step on\-device and synchronises across steps with a grid\-wide barrier \(Section[2\.2](https://arxiv.org/html/2607.10044#S2.SS2)\)\. The host pays launch overhead once per query rather than once per step, keeping latency in a narrow band asB​WBWgrows \(Figure[3](https://arxiv.org/html/2607.10044#S4.F3)\)\. GPU mean latency rises from0\.560\.56ms atK=100K\{=\}100to1\.911\.91ms atK=1000K\{=\}1000, with p95 below2\.792\.79ms and p99 below3\.313\.31ms\. MARISA\-Opt instead expands beams sequentially on a fixed\-size worker pool \(Section[2\.1](https://arxiv.org/html/2607.10044#S2.SS1)\), and its mean latency rises from9\.039\.03ms to46\.3046\.30ms, while its p99 rises from15\.0815\.08ms to76\.7176\.71ms over the same sweep\. The resulting mean speedup grows from16\.3×16\.3\\timesto24\.2×24\.2\\timesasB​W/KBW/Kincreases \(Figure[3](https://arxiv.org/html/2607.10044#S4.F3)\)\. We observe slight dips nearK≈K\{\\approx\}500 and 800, where the on\-device top\-BBsorter changes thread capacity \(Appendix[D](https://arxiv.org/html/2607.10044#A4)\)\.

MARISA\-Int \(Section[2\.1](https://arxiv.org/html/2607.10044#S2.SS1)\) is intractable at billion\-keyword scale \(\>3\{\>\}3h for a singleK=1000K\{=\}1000query on a 10\-query subsample\), so we adopt MARISA\-Opt as the primary CPU comparator; full MARISA\-Int measurements are in Appendix[H](https://arxiv.org/html/2607.10044#A8)\. Maintaining low latency asB​WBWincreases shows that FlashTrie’s optimizations remain effective at scale; next, we break down kernel runtime by phase to identify which components dominate asB​WBWgrows\.

![Refer to caption](https://arxiv.org/html/2607.10044v1/x3.png)Figure 3:Per\-request trie\-search latency and GPU speedup vs\. beam widthB​WBWatb=1b\{=\}1\(\|𝒱\|\|\\mathcal\{V\}\|=2\.2M, 800M\-key trie\)\.\(Left\)Mean \(solid\) and p95 \(dashed\) latency\. CPU mean rises from9\.09\.0ms \(B​W=100BW\{=\}100\) to46\.346\.3ms \(B​W=1000BW\{=\}1000\), with p9565\.565\.5ms and p9976\.776\.7ms at the top end\. GPU mean stays sub\-22ms across the sweep \(0\.550\.55–1\.911\.91ms\), with p95≤2\.79\\leq 2\.79ms and p99≤3\.31\\leq 3\.31ms\.\(Right\)Mean GPU speedup climbs from16×16\\timesto24×24\\times\. Full per\-percentile table: Appendix[F](https://arxiv.org/html/2607.10044#A6)\.
### 4\.3Runtime Breakdown

To profile the latency trend, we instrument FlashTrie’s cooperative kernel and measure GPU runtime across four phases \(Figure[1](https://arxiv.org/html/2607.10044#S2.F1)\): beam expansion, validation, pruning and grid sync\. Figure[4](https://arxiv.org/html/2607.10044#S4.F4)\(left\) shows the dominant cost shifts with beam width: atK=100K\{=\}100, expansion dominates \(≈\\approx66%\) as trie traversal is the bottleneck\. AsB​WBWgrows, expansion remains inexpensive because parallel beams share upper trie nodes\. In contrast, the number of surviving candidates scales withB×KB\{\\times\}K, increasing validation and pruning cost; atB​W=1000BW\{=\}1000, validation becomes dominant \(≈\\approx43% of total runtime\)\. Pruning remains stable at1515–19%19\\%, while sync overhead stays small \(≤6%\\leq 6\\%, decreasing to3\.5%3\.5\\%\), confirming that the persistent kernel avoids per\-step host orchestration with negligible barrier overhead\. Measurement details are in Appendix[D\.4](https://arxiv.org/html/2607.10044#A4.SS4)\. Overall, the bottleneck shifts from trie traversal to candidate processing at larger BW\. We next quantify the system\-level impact using batch\-throughput measurements\.

### 4\.4Throughput Scaling with Batch Size

Latency improvements are operationally meaningful if they convert into higher serving capacity under realistic batching\. In Figure[4](https://arxiv.org/html/2607.10044#S4.F4)\(right\), the throughput curves reflect different batching regimes\. FlashTrie’s throughput rises steeply fromb=1b\{=\}1tob=8b\{=\}8as batching fills more GPU execution resources, peaking at1,6211\{,\}621q/s forB​W=1000BW\{=\}1000\. This is70×70\\timeshigher than MARISA\-Opt at its own peak batch size\. MARISA\-Opt scales only modestly with batch size because its CPU thread pool is quickly saturated: once all workers are occupied, extra queries mostly queue \(Section[2\.1](https://arxiv.org/html/2607.10044#S2.SS1)\)\. FlashTrie benefits more from batching because each query is mapped to a GPU block with many threads, so larger batches better fill the device’s execution resources\. Beyondb=8b\{=\}8, however, thread synchronization and partition overhead start to dominate, so throughput falls\. These throughput gains indicate more effective GPU utilization under batched serving, with higher occupancy and better device saturation as workload increases\.

![Refer to caption](https://arxiv.org/html/2607.10044v1/x4.png)Figure 4:\(Left\) GPU runtime breakdown of FlashTrie’s cooperative beam\-search kernel vs\. beam widthB​WBW\(b=1b\{=\}1, averaged over 13,000 requests\)\. AsB​WBWgrows, the bottleneck shifts from trie traversal \(beam expansion\) to candidate handling \(beam validation\)\. \(Right\) Throughput \(mean q/s\) vs\.B​WBW, per batch size\. FlashTrie peaks atb=8b\{=\}8for allB​WBW\(6,6236\{,\}623q/s atB​W=100BW\{=\}100,1,6211\{,\}621q/s atB​W=1000BW\{=\}1000\), while MARISA\-Opt flattens earlier\. At each system’s own best batch size, FlashTrie delivers32×32\\times–71×71\\timesmore throughput than MARISA\-Opt\.
### 4\.5Retrieval Quality

Given that FlashTrie maintains only a fraction of MARISA\-Opt’s latency even at larger beam widths, we next verify that this highly optimized decoding path does not compromise retrieval quality\. Table[1](https://arxiv.org/html/2607.10044#S4.T1)shows that retrieval quality improves with beam width\. P@100 increases from 0\.53 atB​W=100BW\{=\}100to 0\.78 atB​W=1000BW\{=\}1000, and P@200 increases from 0\.52 to 0\.69\. Gains taper afterB​W≈700BW\{\\approx\}700, but the trend remains monotonic overall at the operating scale of interest\.

Table 1:Retrieval precision at cutoffs 100 and 200 on the 13K\-request dataset, scored by a Transformer cross\-encoderValluriet al\.\([2025](https://arxiv.org/html/2607.10044#bib.bib30)\)\. MARISA\-Opt and FlashTrie agree within0\.0010\.001for everyB​WBW\(collapsed into one column\)\. Prec@200 atB​W=100BW\{=\}100equals Prec@100 due to a truncated denominatorFlashTrie matches MARISA\-Opt to within0\.0010\.001across all settings, indicating that GPU execution remains in parity with the CPU implementation\. More importantly, in latency\-critical serving \(e\.g\., sponsored search; Section[4\.8](https://arxiv.org/html/2607.10044#S4.SS8)\), MARISA\-Opt becomes impractical beyond roughlyB​W≥200BW\{\\geq\}200as beam width grows\. Thus, the quality gains from larger beam widths are operationally realizable only with FlashTrie, which keeps constrained decoding within tight online latency budgets where large\-B​WBWMARISA\-Opt cannot\.

### 4\.6Ablations: Isolating Trie and Search\-Algorithm Contributions

Section[4\.2](https://arxiv.org/html/2607.10044#S4.SS2)combines three effects: execution substrate \(CPU vs\. GPU\), trie\-based constraint checking, and inner\-loop child\-search algorithm\. We isolate the latter two with GPU\-resident ablations that share FlashTrie’s cooperative kernel, scoring rules, thresholds, and length normalization\.

Per\-Position Token \(PPT\) is a depth\-only filter, not a prefix\-constrained decoder: at steptt, it intersects top\-KKmodel tokens with a global sorted token list for depthtt, then returns all matches regardless of beam prefix\. Because this depth\-level token set is shared, the matched set is computed once and reused across beams\.

Linear\-probe keeps the LOUDS trie unchanged but swaps binary search for a shared\-memory linear scan over children, isolating𝒪​\(Δ\)\\mathcal\{O\}\(\\Delta\)vs\.𝒪​\(log⁡Δ\)\\mathcal\{O\}\(\\log\\Delta\)behavior\. Setup details are in Appendices[I](https://arxiv.org/html/2607.10044#A9)and[J](https://arxiv.org/html/2607.10044#A10.SSx2)\.

Table 2:Per\-request latency \(ms\) for two GPU ablations vs\. FlashTrie atb=1b\{=\}1\. PPT removes trie constraints; Linear\-probe keeps the trie but replaces binary search with linear scan in the child range\. We report p95 speedup asablation p95/FlashTrie p95\\text\{ablation p95\}/\\text\{FlashTrie p95\}\. Every otherKKis shown; full sweep in Appendix[J](https://arxiv.org/html/2607.10044#A10), Table[8](https://arxiv.org/html/2607.10044#A10.T8)\.Table[2](https://arxiv.org/html/2607.10044#S4.T2)shows three takeaways\. \(i\) A permissive GPU baseline \(PPT\) is fast but not constraint\-equivalent\. \(ii\) adding prefix\-constrained trie checking \(PPT to FlashTrie\) provides an additional gain that grows withB​WBW\(from1\.2×1\.2\\timesto6\.0×6\.0\\timesin p95\)\. \(iii\) binary search over child labels is critical \(Linear\-probe to FlashTrie\), with71×71\\times–209×209\\timesp95 degradation when replaced by linear scan\. These ablations show that both prefix\-constrained trie filtering and logarithmic child lookup are essential, indicating the gains come from core algorithmic design rather than hardware alone\.

### 4\.7Public\-Dataset Reproducibility

To reproduce the latency claim on public artifacts, we synthesise a parallel workload from NQ\-Open queriesKwiatkowskiet al\.\([2019](https://arxiv.org/html/2607.10044#bib.bib28)\); Leeet al\.\([2019](https://arxiv.org/html/2607.10044#bib.bib32)\), the GENRE\-KILT BART\-large checkpointCaoet al\.\([2021](https://arxiv.org/html/2607.10044#bib.bib2)\); Lewiset al\.\([2020](https://arxiv.org/html/2607.10044#bib.bib33)\), and the∼\\sim6M KILT Wikipedia titlesPetroniet al\.\([2021](https://arxiv.org/html/2607.10044#bib.bib29)\)as the constraint library\. We sweep beam widthB​W∈\{100,200,…,1000\}BW\\in\\\{100,200,\\ldots,1000\\\}over all 3,600 NQ\-Open validation queries \(T=16T\{=\}16decoding positions\) and measure per\-request trie\-search latency on the same hardware as Section[3](https://arxiv.org/html/2607.10044#S3)\.

Figure[5](https://arxiv.org/html/2607.10044#S4.F5)shows that FlashTrie reproduces the internal\-dataset speedup trend \(Figure[3](https://arxiv.org/html/2607.10044#S4.F3)\) on entirely public artifacts\. Thep​99p99speedup increases from12\.8×12\.8\\timesto22\.7×22\.7\\times, indicating both lower and more predictable tail latency across the fullB​WBW\-sweep\. A detailed noise\-injection ablation studying the effect of cross\-position incoherence \(NAR\-style proposal grids\) is in Appendix[K](https://arxiv.org/html/2607.10044#A11)\. Reproducing the same scaling pattern on public data supports generality beyond internal traffic, leading to the final test on live production impact\.

![Refer to caption](https://arxiv.org/html/2607.10044v1/x5.png)Figure 5:Latency on the NQ\+GENRE public workload \(3,600 queries,T=16T\{=\}16,∼6\{\\sim\}6M\-title trie\) vs\. beam widthB​WBW\. GPU mean latency rises from0\.280\.28ms \(B​W=100BW\{=\}100\) to1\.671\.67ms \(B​W=1000BW\{=\}1000\); MARISA\-Opt mean grows from3\.113\.11ms to34\.2334\.23ms\. Mean speedup:11\.0×11\.0\\timesto20\.5×20\.5\\times\.
### 4\.8Online A/B Testing

We validate our offline findings with a randomized online A/B experiment on live traffic from a popular commercial search engine\. The goal is to test whether the larger beam width enabled by FlashTrie yields meaningful online gains in user engagement and revenue, i\.e\., whether offline improvements transfer to gains in production\. The experiment runs for 16 consecutive days across multiple countries, with English and non\-English results reported separately in Table[3](https://arxiv.org/html/2607.10044#S4.T3)\.

Both arms deploy the same NAR generative retrieval modelValluriet al\.\([2025](https://arxiv.org/html/2607.10044#bib.bib30)\)on an A100 MIG 10 GB slice\. In control, production MARISA\-Opt atB​W=200BW\{=\}200is already near the online limit \(≈\\approx25 ms p95;B​W\>200BW\{\>\}200violates latency constraints\)\. The treatment swaps MARISA\-Opt trie stage for FlashTrie and uses the recovered latency headroom to increase beam width toB​Wtrt=600BW\_\{\\text\{trt\}\}\{=\}600, while keeping the model and downstream ranking stack unchanged\. FlashTrie sustainsB​W=600BW\{=\}600at15/1715/17ms \(p50/p95\), exploring over2×2\\timesmore candidates within the same serving budget\.

The treatment yields statistically significant gains, including a0\.71%0\.71\\%lift in revenue, without degradation in ad quality\. We further analyze key metrics ad coverage, impressions, clicks and defect rate\. Ad defect, measured using offline relevance models, denotes the proportion of irrelevant ads shown to users\. FlashTrie increases user engagement, with statistically significant click improvements of0\.17%0\.17\\%for English queries and0\.20%0\.20\\%for non\-English queries\. Since the underlying model and its top\-K proposals are the same, these gains result mainly from increasing beam width, enabling retrieval of more candidate ads and expanding the proposal set passed to downstream ranking\.

Table 3:Online A/B test \(Δ\\Delta\) on a commercial search engine for FlashTrie vs\. CPU MARISA\-Opt, reported by language segment\. All metrics are statistically significant \(α=0\.05\\alpha\{=\}0\.05\)\. See Appendix[G](https://arxiv.org/html/2607.10044#A7)for the full protocol\.Ad coverage and impressions increase as wider beams recover monetizable candidates that narrower beams discard, while revenue gains suggest these candidates correspond to higher\-value matches\. Importantly, defect rates are nonsignificant within±1%\\pm 1\\%, indicating no degradation in result quality\. Overall, the latency headroom enabled by FlashTrie allows a substantial expansion of beam width, translating directly into improved system performance\.

## Discussion

FlashTrie exemplifies a broader accelerator\-era pattern: data structures traditionally implemented on CPUs can become effective on GPUs when their access patterns are redesigned for parallel execution\. Our central finding is that constrained decoding need not remain a CPU bottleneck\. A static succinct trie can be made GPU\-native when lookup and pruning are structured for batched execution and high device utilization\. This reduces decoding latency even at wider beam widths, translating to lower end\-to\-end inference latency under strict serving constraints\.

This shift is important because retrieval quality in generative systems improves at larger beam widths\. In our online setting, FlashTrie enables wider beams within strict latency budgets and converts that headroom into measurable gains in user engagement and revenue without degrading ad quality\. More broadly, FlashTrie applies to constrained generation tasks with large static vocabularies, such as retrieval IDs, product catalogs, or entity linking, that can benefit from accelerator\-resident constraint enforcement that is both fast and scalable\.

## 5Limitations

The current implementation requires the full constraint trie to reside within a single GPU’s VRAM\. At 800M keywords the trie occupies3\.13\.1GB on the A100 80 GB, leaving ample headroom for co\-located weights; for larger libraries a natural extension is to shard subtrees across GPUs over NVLink/NVSwitch and aggregate candidates per step\.

FlashTrie’s batched scheduler allocates a fixed SM partition to each query \(Section[4\.4](https://arxiv.org/html/2607.10044#S4.SS4)\), avoiding cross\-slot synchronisation but capping multi\-query parallelism, dynamic or work\-stealing alternatives that preserve persistent\-kernel coherence are a natural next step\.

All experiments use an NVIDIA A100 80 GB PCIe GPU\. The persistent\-kernel and cooperative\-groups patterns are supported on all CUDA Compute Capability≥7\.0\\geq 7\.0devices, but quantitative characterisation on H100/B100 remains open\.

## References

- Guided open vocabulary image captioning with constrained beam search\.InEMNLP,Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p1.1)\.
- J\. Aoe \(1989\)An efficient digital search algorithm by using a double\-array structure\.IEEE Transactions on Software Engineering15\(9\),pp\. 1066–1077\.External Links:[Document](https://dx.doi.org/10.1109/32.31365)Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p2.1)\.
- M\. Bevilacqua, G\. Ottaviano, P\. Lewis, W\. Yih, S\. Riedel, and F\. Petroni \(2022\)Autoregressive search engines: generating substrings as document identifiers\.InAdvances in Neural Information Processing Systems \(NeurIPS\),External Links:2204\.10628Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p1.1)\.
- N\. D\. Cao, G\. Izacard, S\. Riedel, and F\. Petroni \(2021\)Autoregressive entity retrieval\.InInternational Conference on Learning Representations \(ICLR\),External Links:2010\.00904Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p1.1),[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p2.1),[§4\.7](https://arxiv.org/html/2607.10044#S4.SS7.p1.3)\.
- J\. Daciuk, S\. Mihov, B\. W\. Watson, and R\. E\. Watson \(2000\)Incremental construction of minimal acyclic finite\-state automata\.Computational Linguistics26\(1\),pp\. 3–16\.External Links:[Document](https://dx.doi.org/10.1162/089120100561601)Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p2.1)\.
- T\. Dao, D\. Y\. Fu, S\. Ermon, A\. Rudra, and C\. Ré \(2022\)FlashAttention: fast and memory\-efficient exact attention with IO\-awareness\.InAdvances in Neural Information Processing Systems \(NeurIPS\),External Links:2205\.14135Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p3.1)\.
- J\. Gu, J\. Bradbury, C\. Xiong, V\. O\.K\. Li, and R\. Socher \(2018\)Non\-Autoregressive Neural Machine Translation\.InInternational Conference on Learning Representations,External Links:[Link](https://mlanthology.org/iclr/2018/gu2018iclr-nonautoregressive/)Cited by:[§1](https://arxiv.org/html/2607.10044#S1.p1.1)\.
- C\. Hokamp and Q\. Liu \(2017\)Lexically constrained decoding for sequence to sequence generation using grid beam search\.InProceedings of the 55th Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),Vancouver, Canada,pp\. 1535–1546\.External Links:[Link](https://aclanthology.org/P17-1141)Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p1.1)\.
- G\. Jacobson \(1989\)Succinct static data structures\.Ph\.D\. Thesis,Carnegie Mellon University\.Cited by:[Appendix A](https://arxiv.org/html/2607.10044#A1.p1.1),[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p2.1)\.
- V\. Karpukhin, B\. Oguz, S\. Min, P\. Lewis, L\. Wu, S\. Edunov, D\. Chen, and W\. Yih \(2020\)Dense passage retrieval for open\-domain question answering\.InEMNLP,Cited by:[§1](https://arxiv.org/html/2607.10044#S1.p1.1)\.
- T\. Kwiatkowski, J\. Palomaki, O\. Redfield, M\. Collins, A\. Parikh, C\. Alberti, D\. Epstein, I\. Polosukhin, J\. Devlin, K\. Lee, K\. Toutanova, L\. Jones, M\. Kelcey, M\. Chang, A\. M\. Dai, J\. Uszkoreit, Q\. Le, and S\. Petrov \(2019\)Natural questions: a benchmark for question answering research\.Transactions of the Association for Computational Linguistics7,pp\. 453–466\.External Links:[Document](https://dx.doi.org/10.1162/tacl%5Fa%5F00276)Cited by:[§4\.7](https://arxiv.org/html/2607.10044#S4.SS7.p1.3)\.
- W\. Kwon, Z\. Li, S\. Zhuang, Y\. Sheng, L\. Zheng, C\. H\. Yu, J\. E\. Gonzalez, H\. Zhang, and I\. Stoica \(2023\)Efficient memory management for large language model serving with PagedAttention\.InProceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles \(SOSP\),External Links:2309\.06180Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p3.1)\.
- K\. Lee, M\. Chang, and K\. Toutanova \(2019\)Latent retrieval for weakly supervised open domain question answering\.InACL,pp\. 6086–6096\.Cited by:[§4\.7](https://arxiv.org/html/2607.10044#S4.SS7.p1.3)\.
- M\. Lewis, Y\. Liu, N\. Goyal, M\. Ghazvininejad, A\. Mohamed, O\. Levy, V\. Stoyanov, and L\. Zettlemoyer \(2020\)BART: denoising sequence\-to\-sequence pre\-training for natural language generation, translation, and comprehension\.InACL,pp\. 7871–7880\.Cited by:[§4\.7](https://arxiv.org/html/2607.10044#S4.SS7.p1.3)\.
- X\. Lu, P\. West, R\. Zellers, R\. L\. Bras, C\. Bhagavatula, and Y\. Choi \(2021\)NeuroLogic decoding: \(un\)supervised neural text generation with predicate logic constraints\.InProceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies,External Links:2010\.12884Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p1.1)\.
- D\. Merrill, M\. Garland, and A\. Grimshaw \(2012\)Scalable GPU graph traversal\.InProceedings of the 17th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming \(PPoPP\),pp\. 117–128\.External Links:[Document](https://dx.doi.org/10.1145/2145816.2145832)Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p2.1)\.
- D\. Metzler, Y\. Tay, D\. Bahri, and M\. Najork \(2021\)Rethinking search: making domain experts out of dilettantes\.SIGIR Forum55\(1\),pp\. 1–27\.Cited by:[§1](https://arxiv.org/html/2607.10044#S1.p1.1)\.
- D\. R\. Morrison \(1968\)PATRICIA—practical algorithm to retrieve information coded in alphanumeric\.Journal of the ACM15\(4\),pp\. 514–534\.External Links:[Document](https://dx.doi.org/10.1145/321479.321481)Cited by:[Appendix A](https://arxiv.org/html/2607.10044#A1.SS0.SSS0.Px1.p1.2),[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p2.1)\.
- NVIDIA Corporation \(2020\)NVIDIA a100 tensor core gpu architecture\.Note:[https://resources\.nvidia\.com/en\-us\-tensor\-core/nvidia\-ampere\-architecture\-whitepaper](https://resources.nvidia.com/en-us-tensor-core/nvidia-ampere-architecture-whitepaper)Accessed: 2026\-06\-17Cited by:[Appendix D](https://arxiv.org/html/2607.10044#A4.p1.5),[§2\.2](https://arxiv.org/html/2607.10044#S2.SS2.SSS0.Px2.p1.1)\.
- NVIDIA Corporation \(2024\)CUDA c\+\+ programming guide\.Note:[https://docs\.nvidia\.com/cuda/cuda\-c\-programming\-guide/](https://docs.nvidia.com/cuda/cuda-c-programming-guide/)Accessed: 2026\-06\-17Cited by:[Appendix D](https://arxiv.org/html/2607.10044#A4.p1.5),[§2\.2](https://arxiv.org/html/2607.10044#S2.SS2.SSS0.Px2.p1.1)\.
- G\. Penha, E\. D’Amico, M\. D\. Nadai, E\. Palumbo, A\. Tamborrino, A\. Vardasbi, M\. Lefarov, S\. Lin, T\. Heath, F\. Fabbri, and H\. Bouchard \(2025\)Semantic IDs for joint generative search and recommendation\.arXiv preprint arXiv:2508\.10478\.External Links:2508\.10478Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p1.1)\.
- F\. Petroni, A\. Piktus, A\. Fan, P\. Lewis, M\. Yazdani, N\. De Cao, J\. Thorne, Y\. Jernite, V\. Karpukhin, J\. Maillard, V\. Plachouras, T\. Rocktäschel, and S\. Riedel \(2021\)KILT: a benchmark for knowledge intensive language tasks\.InProceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies,pp\. 2523–2544\.External Links:[Document](https://dx.doi.org/10.18653/v1/2021.naacl-main.200)Cited by:[§4\.7](https://arxiv.org/html/2607.10044#S4.SS7.p1.3)\.
- R\. Pradeep, K\. Hui, J\. Gupta, H\. Zhuang, A\. Lelkes, J\. Lin, D\. Metzler, and V\. Tran \(2023\)Understanding generative retrieval at scale\.InEMNLP,Cited by:[§1](https://arxiv.org/html/2607.10044#S1.p1.1)\.
- Z\. Sun and Y\. Yang \(2020\)An em approach to non\-autoregressive conditional sequence generation\.InICML,Cited by:[§1](https://arxiv.org/html/2607.10044#S1.p1.1)\.
- I\. Sutskever, O\. Vinyals, and Q\. Le \(2014\)Sequence to sequence learning with neural networks\.InNeurIPS,Cited by:[§1](https://arxiv.org/html/2607.10044#S1.p1.1)\.
- Y\. Tay, V\. Q\. Tran, M\. Dehghani, J\. Ni, D\. Bahri, H\. Mehta, Z\. Qin, K\. Hui, Z\. Zhao, J\. Gupta, T\. Schuster, W\. W\. Cohen, and D\. Metzler \(2022\)Transformer memory as a differentiable search index\.InAdvances in Neural Information Processing Systems \(NeurIPS\),External Links:2202\.06991Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p1.1),[§1](https://arxiv.org/html/2607.10044#S1.p1.1)\.
- R\. Valluri, A\. K\. Mohankumar, K\. Dave, A\. Singh, J\. Jiao, M\. Varma, and G\. Sinha \(2025\)Scaling the vocabulary of non\-autoregressive models for fast generative retrieval\.InProceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining V\.1,KDD ’25,New York, NY, USA,pp\. 1409–1420\.External Links:ISBN 9798400712456,[Link](https://doi.org/10.1145/3690624.3709330),[Document](https://dx.doi.org/10.1145/3690624.3709330)Cited by:[Appendix E](https://arxiv.org/html/2607.10044#A5.SS0.SSS0.Px5.p1.3),[§4\.8](https://arxiv.org/html/2607.10044#S4.SS8.p2.7),[Table 1](https://arxiv.org/html/2607.10044#S4.T1)\.
- X\. Wang, Y\. Xiong, Y\. Wei, M\. Wang, and L\. Li \(2021\)LightSeq: a high performance inference library for transformers\.InProceedings of the 2021 Annual Conference of the North American Chapter of the Association for Computational Linguistics: Industry Track \(NAACL\-Industry\),External Links:2010\.13887Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p3.1)\.
- Y\. Wang, Y\. Hou, H\. Wang, Z\. Miao, S\. Wu, H\. Sun, Q\. Chen, Y\. Xia, C\. Chi, G\. Zhao, Z\. Liu, X\. Xie, H\. A\. Sun, W\. Deng, Q\. Zhang, and M\. Yang \(2022\)A neural corpus indexer for document retrieval\.InAdvances in Neural Information Processing Systems \(NeurIPS\),External Links:2206\.02743Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p1.1)\.
- B\. T\. Willard and R\. Louf \(2023\)Efficient guided generation for large language models\.arXiv preprint\.External Links:2307\.09702Cited by:[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p1.1)\.
- Y\. Wu, M\. Schuster, Z\. Chen, Q\. V\. Le, M\. Norouzi, W\. Macherey, M\. Krikun, Y\. Cao, Q\. Gao, K\. Macherey,et al\.\(2016\)Google’s neural machine translation system: bridging the gap between human and machine translation\.arXiv preprint arXiv:1609\.08144\.Cited by:[Appendix E](https://arxiv.org/html/2607.10044#A5.SS0.SSS0.Px2.p1.4),[§2](https://arxiv.org/html/2607.10044#S2.p1.7)\.
- S\. Yata \(2011\)MARISA\-trie: matching algorithm with recursively implemented StorAge\.Note:[https://github\.com/s\-yata/marisa\-trie](https://github.com/s-yata/marisa-trie)Cited by:[Appendix A](https://arxiv.org/html/2607.10044#A1.SS0.SSS0.Px1.p1.2),[Appendix A](https://arxiv.org/html/2607.10044#A1.SS0.SSS0.Px1.p1.3),[Appendix A](https://arxiv.org/html/2607.10044#A1.SS0.SSS0.Px2.p1.1),[Appendix A](https://arxiv.org/html/2607.10044#A1.p1.1),[§B\.1](https://arxiv.org/html/2607.10044#A2.SS1.p1.1),[§1](https://arxiv.org/html/2607.10044#S1.SS0.SSS0.Px1.p2.1),[§2\.1](https://arxiv.org/html/2607.10044#S2.SS1.p1.3),[§2](https://arxiv.org/html/2607.10044#S2.p1.7)\.
- N\. Ziems, W\. Yu, Z\. Zhang, and M\. Jiang \(2023\)Large language models are built\-in autoregressive search engines\.InFindings of the Association for Computational Linguistics \(ACL Findings\),External Links:2305\.09612Cited by:[§1](https://arxiv.org/html/2607.10044#S1.p1.1)\.

Appendix

## Appendix ATrie Background: LOUDS with Tail/Link Compression

All three systems share the same LOUDS backboneYata \([2011](https://arxiv.org/html/2607.10044#bib.bib14)\); Jacobson \([1989](https://arxiv.org/html/2607.10044#bib.bib15)\): a2​n\+12n\{\+\}1\-bit topology vector with constant\-timerank/select, plus per\-nodebases\_labels andterminal\_flags\_\. What differs across systems is not the core topology but auxiliary build\-time structures \(Patricia/TAIL, link flags, and cache\) and the search algorithm used on top of this layout \(Table[4](https://arxiv.org/html/2607.10044#A2.T4)\)\.

#### Patricia single\-child compression and the TAIL buffer\.

Patricia compressionMorrison \([1968](https://arxiv.org/html/2607.10044#bib.bib17)\)folds maximal single\-child chains into a head node plus an offset intotail\_; link nodes store offsets, not labels\. Without suffix sharing, this is usually a loss for wide labels\. MARISA’s trie\-of\-tails variant implements this same idea with a recursive TAIL bufferYata \([2011](https://arxiv.org/html/2607.10044#bib.bib14)\)\. The net benefit depends on suffix deduplication, which can be summarized by:

folded​\(k,w,α\)≈w\+α​k​w\+O​\(log⁡\|tail\_\|\)\.\\text\{folded\}\(k,w,\\alpha\)\\;\\approx\\;w\+\\alpha\\,kw\+O\(\\log\|\\texttt\{tail\\\_\}\|\)\.Folding beats expansion iffα​k​w\+w<k​w\\alpha\\,kw\+w<kw, i\.e\.,

α<1−1k\.\\boxed\{\\;\\alpha\\;<\\;1\-\\tfrac\{1\}\{k\}\\,\.\\;\}\(1\)In byte alphabets this criterion often holds; in wide 32\-bit token alphabets suffix sharing is much weaker and the supporting metadata becomes relatively expensive\. This is why MARISA\-Opt disables Patricia/TAIL in our integer\-token setting \(Section[2\.1](https://arxiv.org/html/2607.10044#S2.SS1)\)\. Recursive TAIL \(num\_tries\) helps only when the same criterion holds at each levelYata \([2011](https://arxiv.org/html/2607.10044#bib.bib14)\); we usenum\_tries=1\.

#### Thelink\_flags\_bit\-vector andcache\_\.

link\_flags\_tags whetherbases\_holds a label or a TAIL offset; its rank index and high\-bit arrays scale with the number of link nodes in MARISA’s succinct\-trie layoutYata \([2011](https://arxiv.org/html/2607.10044#bib.bib14)\)\. When Patricia is disabled, these structures largely collapse\. MARISA’s optionalcache\_prefetch table speeds common top\-of\-trie transitions but adds build cost and index bytesYata \([2011](https://arxiv.org/html/2607.10044#bib.bib14)\)\. These component\-level tradeoffs are summarized in Table[4](https://arxiv.org/html/2607.10044#A2.T4); for full implementation details, seeYata \([2011](https://arxiv.org/html/2607.10044#bib.bib14)\)\.

## Appendix BCPU Baselines: Full Design Details and System Contrast

This appendix gives the full design description for the two CPU baselines summarised in Section[2\.1](https://arxiv.org/html/2607.10044#S2.SS1)and the axis\-by\-axis contrast between all three systems \(Table[4](https://arxiv.org/html/2607.10044#A2.T4)\)\.

### B\.1MARISA\-Int

MARISA\-Int extends the open\-source MARISA codebaseYata \([2011](https://arxiv.org/html/2607.10044#bib.bib14)\)with 32\-bit integer\-token sequences, reusing the original character\-oriented storage layout and traversal routines without modifying the build pipeline\. Since upstream MARISA provides only single\-sequence operations \(e\.g\.,lookup,predictive\_search\), MARISA\-Int adds a minimal beam\-search layer\. Each candidate is evaluated by repeatedly invoking these routines on the full prefix from the root, performing up to two root\-to\-node traversals per \(beam, proposal\) pair\. No trie state is reused across steps, and pruning is implemented viastd::partial\_sortover a flat vector\. The search is single\-threaded, and all operations execute on the host CPU\.

### B\.2MARISA\-Opt

MARISA\-Opt builds on MARISA\-Int while retaining the LOUDS\-based trie layout and on\-disk format, but introduces a set of implementation optimizations that make it a competitive CPU baseline\. These include compact link\-offset encoding, stateful traversal, multi\-threaded execution, efficient per\-beam expansion, and incremental top\-BBpruning \(full per\-operation derivations in Appendix[C](https://arxiv.org/html/2607.10044#A3)\)\. Comparing FlashTrie against this optimized baseline ensures that performance gains reflect on\-device parallelism rather than unoptimised CPU execution\.

MARISA\-Opt reduces index size by splitting link offsets across the per\-node 32\-bit slot and a secondary array, storing only the low\-order 8 bits in the slot and packing the remaining bits compactly\. The secondary array is allocated only when needed, avoiding unnecessary overhead\. Token labels themselves remain in full in the per\-node slot; only the link offsets are split\. On CPU, MARISA\-Opt does not narrow the per\-node slot itself, the slot is still 32 bits, so at each link node 24 of those 32 bits are now unused padding\.

For integer\-token alphabets, MARISA\-Opt further removes structures whose benefit is specific to byte\-alphabet inputs \(Appendix[A](https://arxiv.org/html/2607.10044#A1)\)\. First, Patricia single\-child link compression is disabled in the build path: every input token becomes its own LOUDS node, no link nodes are emitted, and the recursive trie\-of\-tails terminates after the first level regardless of the configurednum\_tries\. As a consequence,extras\_, the TAIL buffer, and the rank/select index overlink\_flags\_are all empty in the on\-disk index, and the link\-offset split encoding described above is exercised only by character workloads\. Second, the build\-time prefetch cache table \(cache\_, used to amortize top\-of\-trie descents at lookup time\) is skipped at construction and is absent from the on\-disk index\. Third, the in\-process multikey quicksort overVector<Key\>is replaced by an externalsort \-u \-VUnix pipeline that delivers unique, lexicographically ordered token sequences to the builder, so duplicate keys are removed before any trie work begins\.

The combined effect is that on a 32\-bit alphabet, where each TAIL token would still occupy 4 bytes and each TAIL reference would carry non\-trivialextras\_andlink\_flags\_overhead, the bytes that Patricia compression would remove frombases\_are smaller than the bytes its supporting structures would add\. Disabling the compression therefore shrinks the on\-disk index and shortens construction simultaneously; this is the dominant driver of MARISA\-Opt’s index\-size and build\-time advantage over MARISA\-Int in Section[4\.1](https://arxiv.org/html/2607.10044#S4.SS1)\.

Unlike MARISA\-Int, which performs repeated root\-to\-node traversals at each decoding step, MARISA\-Opt propagates trie state across beams\. Each extension requires a single LOUDS navigation step from the cached node, making per\-candidate cost independent of prefix length\.

Work is parallelized across beams using a fixed\-size CPU worker pool, while per\-beam expansion is performed using a linear\-time two\-pointer merge between sorted proposals and sorted trie children\. This reduces matching cost fromO​\(K​Δ\)O\(K\\Delta\)toO​\(K\+Δ\)O\(K\{\+\}\\Delta\)with sequential memory access\. Threshold\-based pruning is applied during expansion to discard low\-probability candidates early\.

Instead of end\-of\-step sorting, MARISA\-Opt maintains a shared bounded min\-heap of sizeB​WBW, reducing pruning cost fromO​\(C​log⁡C\)O\(C\\log C\)toO​\(C​log⁡B​W\)O\(C\\log BW\), whereCCis the number of candidates generated per step andB​WBWis the beam width, while keeping the top\-BBbeams continuously updated\.

### B\.3Axis\-by\-Axis System Contrast

Table 4:Design and implementation contrast across the three systems\. All three share the LOUDS\+\+tail/link trie skeleton \(Appendix[A](https://arxiv.org/html/2607.10044#A1)\); they differ in build\-time storage choices and in where/how the beam search is executed\. The*Improves*column tags whether each axis primarily affects search latency \(LAT\), build time \(BUILD\), or trie\-on\-disk size \(SIZE\)\. Per\-operation time complexity is reported separately in Table[5](https://arxiv.org/html/2607.10044#A4.T5)\. Background on Patricia single\-child compression, the TAIL buffer, thelink\_flags\_bit\-vector, and the prefetchcache\_table referenced in the rows below is given in Appendix[A](https://arxiv.org/html/2607.10044#A1)\.

## Appendix CCPU Trie Operations

The CPU trie operations used by MARISA\-Int and MARISA\-Opt are inherited from the MARISA framework and are not contributions of this work; we describe them in prose for completeness so that the GPU adaptations in Appendix[D](https://arxiv.org/html/2607.10044#A4)have a precise reference point\.

#### Child enumeration\.

Locating the children of a nodevvis two rank/select operations onlouds\_\. The position ofvv’s first child withinlouds\_isselect0⁡\(v\)\+1\\operatorname\{select\}\_\{0\}\(v\)\+1, and the corresponding child node index isselect0⁡\(v\)\+1−v−1\\operatorname\{select\}\_\{0\}\(v\)\+1\-v\-1\(subtracting the LOUDS header bits seen so far\)\. The number of childrenΔ\\Deltais the distance from that position to the next0\-bit; on CPU this is a sequential scan overlouds\_, and on GPU it is replaced byWarpNextUnset\(Appendix[D](https://arxiv.org/html/2607.10044#A4)\)\. Labels of theΔ\\Deltachildren sit contiguously atbases\_\[base\_node\.\.base\_node\+Δ\\Delta\)in sorted order\.

#### FindChild\.

Given a parentvvand a target tokenτ\\tau, MARISA does a binary search over theΔ\\Deltachild labels\. The only subtlety is that when a childuuis a link node \(link\_flags\_​\[u\]=1\\texttt\{link\\\_flags\\\_\}\[u\]=1\), itsbases\_slot holds a TAIL offset rather than a token, so the binary\-search comparison must peek intotail\_at that offset to recover the first label of the suffix; the comparison itself is then ordinary integer comparison\. This costs one extra memory indirection per visited link node and is the reason TAIL traversal must read at most one tail byte per binary\-search step rather than the entire suffix\. Per\-call cost isO​\(log⁡Δ\)O\(\\log\\Delta\)comparisons, eachO​\(1\)O\(1\)amortized\.

#### Lookupand predictive search\.

A whole\-key lookup isLLsuccessiveFindChildcalls followed by a check that the resulting node carries a terminal flag; the key identifier is thenrank1⁡\(terminal\_flags\_,v\)−1\\operatorname\{rank\}\_\{1\}\(\\texttt\{terminal\\\_flags\\\_\},v\)\-1\. Common prefix search and predictive search extend this loop with an additional in\-order walk of the subtree rooted at the deepest matching node, payingO​\(R\)O\(R\)forRRreturned keys on top of theO​\(L​log⁡Δ\)O\(L\\log\\Delta\)descent\. All of these operations execute serially on CPU and are the per\-beam unit of work that FlashTrie’s expansion kernel parallelises in Appendix[D](https://arxiv.org/html/2607.10044#A4)\.

## Appendix DGPU Implementation Details

FlashTrie replaces the serialFindChild/Lookuploop of Appendix[C](https://arxiv.org/html/2607.10044#A3)with a cooperative\-kernel pipeline that keeps the trie, the beam state, and the top\-KKoutput resident on the device for the entire queryNVIDIA Corporation \([2024](https://arxiv.org/html/2607.10044#bib.bib3),[2020](https://arxiv.org/html/2607.10044#bib.bib4)\)\. This appendix documents the pieces that genuinely differ from the CPU baseline and are not adequately specified by prose alone: tail\-suffix traversal, warp\-cooperative LOUDS scanning, and large\-B​WBWtop\-KKselection\. Bookkeeping aspects, memory residency, output serialization, and per\-batch allocation, are described in prose since they involve no concurrency subtleties\. Table[5](https://arxiv.org/html/2607.10044#A4.T5)summarises per\-operation time complexity for the GPU primitives used below alongside the CPU counterparts; the1/P1/Pfactor is the per\-thread reduction the GPU contributes, withP=512P\{=\}512for per\-CTA expansion and the full grid for sort passes\.

Table 5:Per\-operation time complexity\.LL= key length,QQ= query length,RR= result count,Δ\\Delta= node out\-degree,TT= decoding steps,B​WBW= beam width,B​WBW= top\-KKproposals per step,PP= GPU parallelism, BS = Beam Search\.#### Memory architecture\.

GPU memory is partitioned into three classes\. The*static trie*\(LOUDS bit\-vectors, label arrays, tail buffer, rank/select indices\) lives in CUDA Unified Memory and is prefetched on demand; this lets a single index back many concurrent search slots without per\-slot duplication\.*Per\-batch search state*, three frontier arenas, a backtrace arena for parent pointers, and a beam\-history buffer used during output reconstruction, is device\-resident, allocated once atinit\_tbs\_gputime, and reused across requests so that nocudaMallocappears on the hot path\.*Per\-request I/O*\(input top\-KKtokens and log\-probabilities, output beam sequences and normalised scores\) is held in pinned host buffers paired with device mirrors, enabling asynchronous DMA that overlaps with kernel execution\.

### D\.1Tail Traversal

Beam expansion on the GPU runs in two passes per decoding step, both driven by the cooperative kernel in Algorithm[1](https://arxiv.org/html/2607.10044#alg1)\. The first pass handles beams that are already*inside*a tail suffix \(s\.in\_link\): such a beam has exactly one valid continuation, the next token along the tail chain, so the work is to \(i\) read that token fromtail\_ats\.link\_offset, \(ii\) scan the top\-KKproposals for a match, \(iii\) apply token\- and sentence\-level log\-prob thresholds, and \(iv\) emit the resulting beam either toselected\(if the proposal closes the chain on a terminal node\) or tonext\. Because the continuation is unique, no binary search is needed and each beam costsO​\(K/P\)O\(K/P\)work across the CTA’sPPthreads; this isExtendInLinkin Algorithm[2](https://arxiv.org/html/2607.10044#alg2)\.

The second pass handles beams sitting on ordinary LOUDS nodes\. For each such beam the kernel first computes the child range\[𝑏𝑎𝑠𝑒​\_​𝑛𝑜𝑑𝑒,𝑏𝑎𝑠𝑒​\_​𝑛𝑜𝑑𝑒\+Δ\)\[\\mathit\{base\\\_node\},\\mathit\{base\\\_node\}\{\+\}\\Delta\)usingselect0\\operatorname\{select\}\_\{0\}andWarpNextUnset, then has each thread test one top\-KKproposal in parallel by binary\-searching the sorted child labels\. Each comparison transparently peeks intotail\_when the visited child is a link node, exactly as in CPUFindChild\. Surviving proposals become new beams that are routed toselectedornexton the same terminal\-/link\-status rules as the first pass; this isExtendNodein Algorithm[3](https://arxiv.org/html/2607.10044#alg3)\. Splitting the work into two passes keeps the warps that are walking tail chains from interfering with the warps that are doing wider parallel child search, which keeps lane utilisation high in both modes\.

Algorithm 2ExtendInLink \(device; one CTA per parent beam\)1:

τnext,𝑒𝑛𝑑←𝑡𝑎𝑖𝑙\.next​\_​token\(s\.𝑙𝑖𝑛𝑘\_𝑜𝑓𝑓𝑠𝑒𝑡\)\\tau\_\{\\mathrm\{next\}\},\\,\\mathit\{end\}\\leftarrow\\mathit\{tail\}\.\\operatorname\{next\\\_token\}\(s\.\\mathit\{link\\\_offset\}\)
2:for

k=𝑡ℎ𝑟𝑒𝑎𝑑𝐼𝑑𝑥\.xk=\\mathit\{threadIdx\.x\}to

K−1K\-1step

𝑏𝑙𝑜𝑐𝑘𝐷𝑖𝑚\.x\\mathit\{blockDim\.x\}do

3:if

𝑡𝑜𝑝𝑘​\_​𝑖𝑑​\[k\]≠τnext\\mathit\{topk\\\_id\}\[k\]\\neq\\tau\_\{\\mathrm\{next\}\}then

4:continue

5:endif

6:

ρ←𝑡𝑜𝑝𝑘​\_​𝑙𝑜𝑔𝑝​\[k\]\\rho\\leftarrow\\mathit\{topk\\\_logp\}\[k\]
7:if

ρ≤θtok\\rho\\leq\\theta\_\{\\mathrm\{tok\}\}or

s\.σ\+ρ≤θsents\.\\sigma\+\\rho\\leq\\theta\_\{\\mathrm\{sent\}\}then

8:break

9:endif

10:

s′←BeamState​\(𝑝𝑎𝑟𝑒𝑛𝑡=s,𝑡𝑜𝑘𝑒𝑛=τnext\)s^\{\\prime\}\\leftarrow\\textsc\{BeamState\}\(\\mathit\{parent\}\{=\}s,\\,\\mathit\{token\}\{=\}\\tau\_\{\\mathrm\{next\}\}\)
11:

s′\.σ←s\.σ\+ρs^\{\\prime\}\.\\sigma\\leftarrow s\.\\sigma\+\\rho
12:

s′\.𝑙𝑖𝑛𝑘​\_​𝑜𝑓𝑓𝑠𝑒𝑡←s\.𝑙𝑖𝑛𝑘​\_​𝑜𝑓𝑓𝑠𝑒𝑡\+1s^\{\\prime\}\.\\mathit\{link\\\_offset\}\\leftarrow s\.\\mathit\{link\\\_offset\}\+1;

s′\.𝑒𝑛𝑑​\_​𝑙𝑖𝑛𝑘←𝑒𝑛𝑑s^\{\\prime\}\.\\mathit\{end\\\_link\}\\leftarrow\\mathit\{end\}
13:if

𝑒𝑛𝑑\\mathit\{end\}then

14:if

is​\_​terminal\(s′\.v\)\\operatorname\{is\\\_terminal\}\(s^\{\\prime\}\.v\)then

15:

𝑠𝑒𝑙𝑒𝑐𝑡𝑒𝑑\.push⁡\(s′\)\\mathit\{selected\}\.\\operatorname\{push\}\(s^\{\\prime\}\)
16:endif

17:else

18:

𝑛𝑒𝑥𝑡\.push⁡\(s′\)\\mathit\{next\}\.\\operatorname\{push\}\(s^\{\\prime\}\)
19:endif

20:break

21:endfor

Algorithm 3ExtendNode \(device; one CTA per parent beam\)1:Compute child range

\[𝑏𝑎𝑠𝑒​\_​𝑛𝑜𝑑𝑒,𝑏𝑎𝑠𝑒​\_​𝑛𝑜𝑑𝑒\+𝑑𝑒𝑔\)\[\\mathit\{base\\\_node\},\\,\\mathit\{base\\\_node\}\+\\mathit\{deg\}\)via LOUDS

select0\\operatorname\{select\}\_\{0\}andWarpNextUnset\(\)

2:if

𝑑𝑒𝑔=0\\mathit\{deg\}=0then

3:return

4:endif

5:for

k=𝑡ℎ𝑟𝑒𝑎𝑑𝐼𝑑𝑥\.xk=\\mathit\{threadIdx\.x\}to

K−1K\-1step

𝑏𝑙𝑜𝑐𝑘𝐷𝑖𝑚\.x\\mathit\{blockDim\.x\}do

6:

τ←𝑡𝑜𝑝𝑘​\_​𝑖𝑑​\[k\]\\tau\\leftarrow\\mathit\{topk\\\_id\}\[k\];

ρ←𝑡𝑜𝑝𝑘​\_​𝑙𝑜𝑔𝑝​\[k\]\\rho\\leftarrow\\mathit\{topk\\\_logp\}\[k\]
7:if

ρ≤θtok\\rho\\leq\\theta\_\{\\mathrm\{tok\}\}or

s\.σ\+ρ≤θsents\.\\sigma\+\\rho\\leq\\theta\_\{\\mathrm\{sent\}\}then

8:continue

9:endif

10:Binary\-search sorted child labels in

\[𝑏𝑎𝑠𝑒​\_​𝑛𝑜𝑑𝑒,𝑏𝑎𝑠𝑒​\_​𝑛𝑜𝑑𝑒\+𝑑𝑒𝑔\)\[\\mathit\{base\\\_node\},\\,\\mathit\{base\\\_node\}\+\\mathit\{deg\}\)for

τ\\tau; peek into

𝑡𝑎𝑖𝑙\\mathit\{tail\}when the visited child is a link node

11:ifnot foundthen

12:continue

13:endif

14:

s′←BeamState​\(𝑝𝑎𝑟𝑒𝑛𝑡=s,𝑡𝑜𝑘𝑒𝑛=τ\)s^\{\\prime\}\\leftarrow\\textsc\{BeamState\}\(\\mathit\{parent\}\{=\}s,\\,\\mathit\{token\}\{=\}\\tau\);

s′\.σ←s\.σ\+ρs^\{\\prime\}\.\\sigma\\leftarrow s\.\\sigma\+\\rho;

s′\.v←s^\{\\prime\}\.v\\leftarrowmatched child

15:Route

s′s^\{\\prime\}to

𝑠𝑒𝑙𝑒𝑐𝑡𝑒𝑑\\mathit\{selected\}or

𝑛𝑒𝑥𝑡\\mathit\{next\}per terminal\- and link\-status \(same rules as Algorithm[2](https://arxiv.org/html/2607.10044#alg2)\)

16:endfor

### D\.2Warp\-Level LOUDS Navigation

The inner step of child enumeration, finding the next0\-bit inlouds\_starting from a given position, can span many 64\-bit words when a node has high out\-degree, and a sequential per\-thread scan leaves 31 of the 32 lanes of a warp idle\. FlashTrie uses a warp\-ballot loop in which each lane checks one of the next 32 words simultaneously and the warp votes via\_\_ballot\_sync; the position of the first word containing a0\-bit is recovered byffson the ballot mask, and the final intra\-word bit position by anotherffson the negated word\. Asymptotically this turns theO​\(Δ\)O\(\\Delta\)inner loop intoO​\(Δ/32\)O\(\\Delta/32\)memory accesses\. Algorithm[4](https://arxiv.org/html/2607.10044#alg4)gives the exact word\-level masking and ballot pattern; the lane arithmetic is what prose cannot specify precisely\.

Algorithm 4WarpNextUnset \(device; 32\-thread warp\)1:

w←⌊i/64⌋w\\leftarrow\\lfloor i/64\\rfloor;

𝑤𝑜𝑟𝑑←B\[w\]∣∼\(∼0≪\(imod64\)\)\\mathit\{word\}\\leftarrow B\[w\]\\mid\{\\sim\}\(\{\\sim\}0\\ll\(i\\bmod 64\)\)
2:\{mask bits below

ii\}

3:if

\(∼𝑤𝑜𝑟𝑑\)≠0\(\{\\sim\}\\mathit\{word\}\)\\neq 0then

4:return

w⋅64\+ffs\(∼𝑤𝑜𝑟𝑑\)−1w\\cdot 64\+\\operatorname\{ffs\}\(\{\\sim\}\\mathit\{word\}\)\-1
5:endif

6:

w←w\+1w\\leftarrow w\+1;

𝑙𝑎𝑛𝑒←𝑡ℎ𝑟𝑒𝑎𝑑𝐼𝑑𝑥\.xmod32\\mathit\{lane\}\\leftarrow\\mathit\{threadIdx\.x\}\\bmod 32
7:repeat

8:

𝑏𝑎𝑙𝑙𝑜𝑡←ballot\(\(∼B\[w\+𝑙𝑎𝑛𝑒\]\)≠0\)\\mathit\{ballot\}\\leftarrow\\operatorname\{ballot\}\(\(\{\\sim\}B\[w\+\\mathit\{lane\}\]\)\\neq 0\)
9:if

𝑏𝑎𝑙𝑙𝑜𝑡=0\\mathit\{ballot\}=0then

10:

w←w\+32w\\leftarrow w\+32
11:else

12:

w←w\+ffs⁡\(𝑏𝑎𝑙𝑙𝑜𝑡\)−1w\\leftarrow w\+\\operatorname\{ffs\}\(\\mathit\{ballot\}\)\-1;break

13:endif

14:untilfalse

15:return

w⋅64\+ffs\(∼B\[w\]\)−1w\\cdot 64\+\\operatorname\{ffs\}\(\{\\sim\}B\[w\]\)\-1

### D\.3Top\-KKSelection on the GPU

After expansion the candidate buffernextis trimmed to the top\-BBhypotheses by length\-normalized score\. FlashTrie switches between two implementations depending onB​WBW\.

#### Block merge sort \(B≤1024B\\leq 1024\)\.

For smallB​WBWthe entire selection fits in a small number of sort\-blocks of size2​B2B\. Each CTA sorts one block in descending order ofσ^\\hat\{\\sigma\}using CUB’sBlockMergeSort, then the sorted blocks are combined by a tournament\-merge tree in⌈log2⁡\(n/2​B\)⌉\\lceil\\log\_\{2\}\(n/2B\)\\rceilrounds with a grid\-wide barrier between rounds\. After the final merge round the firstB​WBWentries are the top\-BBbeams\. Total work isO​\(n​log⁡n/P\)O\(n\\log n/P\)for the initial sort andO​\(B​log⁡\(n/B\)/P\)O\(B\\log\(n/B\)/P\)for the merge rounds, withPPthe grid\-wide thread count; see Algorithm[5](https://arxiv.org/html/2607.10044#alg5)\.

Algorithm 5BlockMergeSort Top\-KK\(device; all blocks\)1:

𝑟𝑒𝑚𝑎𝑖𝑛𝑖𝑛𝑔←⌈\|𝑐𝑎𝑛𝑑\|/\(2​B\)⌉\\mathit\{remaining\}\\leftarrow\\lceil\|\\mathit\{cand\}\|/\(2B\)\\rceil
2:\{number of sort\-blocks\}

3:while

𝑟𝑒𝑚𝑎𝑖𝑛𝑖𝑛𝑔≥1\\mathit\{remaining\}\\geq 1do

4:\{Phase 1: sort each block of

2​B2Belements\}

5:for allblock

bbassigned to this CTAdo

6:Load

2​B/P2B/Pelements per thread \(pad with

−∞\-\\inftyif OOB\)

7:

BlockMergeSort​\(𝑡ℎ𝑟𝑒𝑎𝑑​\_​𝑑𝑎𝑡𝑎\)\\textsc\{BlockMergeSort\}\(\\mathit\{thread\\\_data\}\)\(descending by

σ^\\hat\{\\sigma\}\)

8:write back to

𝑠𝑜𝑟𝑡​\_​𝑠𝑡𝑎𝑡𝑒\\mathit\{sort\\\_state\}
9:endfor

10:if

𝑟𝑒𝑚𝑎𝑖𝑛𝑖𝑛𝑔=1\\mathit\{remaining\}=1then

11:break

12:endif

13:

sync⁡\(𝒢\)\\operatorname\{sync\}\(\\mathcal\{G\}\)
14:\{Phase 2: tournament merge\}

15:

𝑟𝑒𝑚𝑎𝑖𝑛𝑖𝑛𝑔←⌈\(𝑟𝑒𝑚𝑎𝑖𝑛𝑖𝑛𝑔−1\)/2⌉\+1\\mathit\{remaining\}\\leftarrow\\lceil\(\\mathit\{remaining\}\-1\)/2\\rceil\+1
16:

sync⁡\(𝒢\)\\operatorname\{sync\}\(\\mathcal\{G\}\)
17:endwhile

18:\{Phase 3: collect top\-

BBentries into

𝑐𝑎𝑛𝑑​\[0​…​B−1\]\\mathit\{cand\}\[0\\ldots B\-1\]\}

#### Output serialization\.

The finalselectedbuffer is descending\-sorted byσ^\\hat\{\\sigma\}but the wire format expects ascending order, soGenerateResultwrites theii\-th selected beam into slotno−1−in\_\{o\}\{\-\}1\{\-\}iof the output arrays\. Per\-beam token sequences are reconstructed by walking thebeam\_historyparent pointers in reverse, with a CUBBlockScanconverting per\-beam lengths into the end\-offsets that index the flattokensbuffer\. The on\-wire layout fornon\_\{o\}beams totallingntokn\_\{\\mathrm\{tok\}\}tokens is an 8\-bytenum\_beamsheader followed by three contiguous arrays:8​no8n\_\{o\}bytes oflogp\_norm\(float64\),8​no8n\_\{o\}bytes ofsent\_offset\(uint64\), and4​ntok4n\_\{\\mathrm\{tok\}\}bytes oftokens\(uint32\); sequenceiioccupiestokens\[sent\_offset\[i−1i\{\-\}1\] : sent\_offset\[ii\]\]\.

#### Per\-batch GPU memory budget\.

A single search slot occupies roughly320320MB of GPU memory, dominated by the persistent device\-resident state: three frontier arrays at5656MB each, a5656MB beam history, six88MB sort\-state arrays, a∼\\sim128128MB device\-side output buffer paired with pinned host memory, and a∼\\sim1212MB input buffer similarly paired\. Top\-KKscratch is small \(3×0\.53\\times 0\.5MB\) and the trie struct shell is a few hundred bytes; the trie arrays themselves live in Unified Memory and range from 50 to 500 MB depending on the constraint set, charged once across all slots rather than per\-slot\. A four\-slot batch therefore uses≈1\.3\\approx 1\.3GB plus the trie, which fits comfortably alongside a generation\-model weights footprint on a single 80 GB device\.

### D\.4Runtime Breakdown Methodology

Because beam search runs as a single persistent cooperative kernel \(Section[2\.2](https://arxiv.org/html/2607.10044#S2.SS2)\), we measure its internal phase breakdown with lightweight on\-device instrumentation rather than separate kernel launches\. Under a compile\-time flag, a single grid\-leader thread reads the GPU cycle counter \(clock64\) along its timeline\. We time the*wall\-clock span*of each phase region, bracketed by the grid barriers \(grid\.sync\) that already separate the phases; because each region ends at a barrier, its span includes the slowest CTA, i\.e\. the phase’s critical\-path time across the grid rather than the leader’s own work alone\. Cycle counts are converted to milliseconds via the SM clock rate and averaged over the 13,000\-request workload\.

The four buckets map to kernel regions as follows\.*Expansion*is the per\-proposal binary search that locates a candidate token among a parent’s children, including the dependent LOUDS/tail memory accesses at each search step \(a single tail read for link\-interior beams\)\.*Validation*is per\-candidate admission: the per\-token and cumulative log\-probability threshold tests, the terminal\-node check, and the append into the candidate buffer\.*Pruning*is the top\-BBselection \(parallel mergesort\) that enforces the beam width\.*Sync*is the residual: inter\-phase grid barriers outside the work regions, frontier setup/save, result serialization, and buffer swaps\. The expansion and validation regions are interleaved per proposal, so we report their combined wall span split by the leader thread’s measured compute ratio; pruning and sync are measured directly\. The four phases sum to total kernel time by construction\.

We report*proportions*rather than absolute times: the instrumented build incurs mild register pressure that inflates absolute latency, so all latency numbers in Section[4\.2](https://arxiv.org/html/2607.10044#S4.SS2)come from a separate, non\-instrumented build\. For reference, the measured mean per\-phase times atb=1b\{=\}1range from0\.340\.34/0\.070\.07/0\.080\.08/0\.030\.03ms \(expansion/validation/pruning/sync\) atB​W=100BW\{=\}100to0\.660\.66/0\.790\.79/0\.350\.35/0\.070\.07ms atB​W=1000BW\{=\}1000, with total kernel time under22ms across the full sweep\.

## Appendix EExperimental Setup: Full Protocol

This appendix fills in the measurement details abbreviated in the main body’s Experimental Setup \(Section[3](https://arxiv.org/html/2607.10044#S3)/Section[3](https://arxiv.org/html/2607.10044#S3)\)\. The deployment envelope \(1×\\timesA100 80 GB, AMD EPYC 7V13 host, 60 GB host RAM, NUMA\-pinned to cores 0–7 of socket 0 vianumactl \-\-cpunodebind=0 \-\-membind=0 taskset \-c 0\-7, Docker container\) and the latency\-window definition \(C\+\+ entry\-point timer covering H2D, kernel, and D2H with stream synchronisation\) are reused from the main body without modification\.

#### Comparability of the timing window\.

The CPU systems’ timer closes when the search call returns; the FlashTrie timer closes aftercudaStreamSynchronize\. All three systems share the same C\+\+ entry\-point boundary and exclude only language\-binding marshalling, so reported latencies are directly comparable\. NUMA\-pinning to a single socket also removes cross\-socket traffic on the H2D path, so FlashTrie’s speedup is reported under the same deployment\-realistic envelope as the CPU baselines\.

#### Length\-normalised scoring\.

Beams are ranked by the Google\-style length\-normalised log\-probability used in Section[2](https://arxiv.org/html/2607.10044#S2)Wuet al\.\([2016](https://arxiv.org/html/2607.10044#bib.bib13)\):

snorm​\(𝐱0:d\)=\(∑j=0d−1ℓj,σ​\(j\)\)⋅\(65\+d\)α,s\_\{\\text\{norm\}\}\(\\mathbf\{x\}\_\{0:d\}\)=\\Bigl\(\\sum\_\{j=0\}^\{d\-1\}\\ell\_\{j,\\sigma\(j\)\}\\Bigr\)\\cdot\\Bigl\(\\frac\{6\}\{5\+d\}\\Bigr\)^\{\\alpha\},\(2\)whereα\\alphacontrols the normalisation strength andσ​\(j\)\\sigma\(j\)indexes the selected candidate at stepjj\.

#### Per\-request latency sampling\.

Each\(B​W,b\)\(BW,b\)configuration is benchmarked over the full 13,000\-request dataset at batch size 1\. One warm\-up pass is discarded; FlashTrie additionally runs a short GPU warm\-up before the measurement window opens\. The next 10 passes are timed and pooled, producing130,000130\{,\}000samples per configuration from which the mean and the 50th/90th/95th/99th percentiles are computed\. The across\-run standard deviation is reported as a variance bar in figures and as the±\\pmvalue in Table[6](https://arxiv.org/html/2607.10044#A6.T6)\.

#### Throughput sweep\.

For the throughput evaluation \(Figure[4](https://arxiv.org/html/2607.10044#S4.F4), Section[4](https://arxiv.org/html/2607.10044#S4)\) we sweepb∈\{1,4,8\}b\\in\\\{1,4,8\\\}atB​W∈\{100,300,500,700,1000\}BW\\in\\\{100,300,500,700,1000\\\}\. For each\(K,b\)\(K,b\)the dataset is run as⌈13000/b⌉\\lceil 13000/b\\rceilbatched calls per pass, and per\-pass throughput is total queries divided by total wall\-clock duration measured with the same C\+\+ entry\-point timer window used for latency\. One warm\-up pass is discarded and four passes are timed; reportedΘ\\Thetais the mean over those four passes, with the across\-pass standard deviation drawn as the error band in Figure[4](https://arxiv.org/html/2607.10044#S4.F4)\.

#### Precision@K scoring\.

For the retrieval\-quality evaluation \(Section[4\.5](https://arxiv.org/html/2607.10044#S4.SS5)\) each request’s top\-KKtrie\-admitted output is scored by a transformer teacher model adapted fromValluriet al\.\([2025](https://arxiv.org/html/2607.10044#bib.bib30)\)that takes the request and a candidate keyword sequence as input and returns a single relevance score\. The teacher is identical across all three systems \(MARISA\-Opt, FlashTrie, and PPT\) so that any precision differences are attributable to the candidate set, not to the scorer\. Precision@B​W=100BW\{=\}100and Precision@B​W=200BW\{=\}200are reported as the macro\-average over the 13,000 requests\.

## Appendix FPer\-Percentile Latency Table

Table[6](https://arxiv.org/html/2607.10044#A6.T6)gives the full per\-percentile breakdown of the per\-request trie\-search latency reported in Figure[3](https://arxiv.org/html/2607.10044#S4.F3): mean, p50, p90, p95 and p99 in milliseconds, for every beam widthB​WBW, computed as the median across 10 runs of 13,000 queries each\.

Table 6:Per\-request trie\-search latency \(ms\), full percentile breakdown\. Each cell is the median across 10 runs of 13,000 queries each\.Lower is better\.
## Appendix GOnline A/B Test Protocol

This appendix gives the operational protocol of the production A/B test summarised in Section[4\.8](https://arxiv.org/html/2607.10044#S4.SS8)\.

#### Deployment motivation\.

Offline benchmarks \(Sections[4\.2](https://arxiv.org/html/2607.10044#S4.SS2)–[4\.5](https://arxiv.org/html/2607.10044#S4.SS5)\) establish that FlashTrie preserves retrieval quality at a fraction of the latency budget of CPU\-resident MARISA\-Opt\. The operative question for a monetised production system is whether these system\-level wins survive the realities of live traffic: bursty query mixes, stale caches, downstream auction dynamics, and revenue\-sensitive guardrails\. The flight is designed to answer that question end\-to\-end\.

#### Sample\-ratio\-mismatch \(SRM\) check\.

A sample\-ratio\-mismatch check passed throughout the flight, confirming the realised treatment : control event ratio is consistent with the configured allocation after eligibility filtering\. This rules out exposure\-bias confounds before reading the metric lifts\.

#### Latency\-timing boundary\.

End\-to\-end latency is measured from request ingress at the retrieval service to emission of the final top\-kkcandidate set, inclusive of trie traversal, beam expansion, scoring, and post\-processing\. Control and treatment share the entire downstream stack \(auction, pacing, allocation, click models\); they differ only in the retrieval\-stage trie back\-end and the beam configuration the back\-end can sustain within the3030ms end\-to\-end latency SLA\.

#### Constraint library and back\-end configurations\.

The production constraint library indexes hundreds of millions of docIDs\. The control arm runs production CPU MARISA\-Opt atBctrl=200B\_\{\\text\{ctrl\}\}\{=\}200,top\-​k=300\\text\{top\-\}k\{=\}300, with a 3\-thread pool matched to the per\-worker MIG\-slice CPU\-core allocation\.

#### Metric glossary\.

Headers in Table[3](https://arxiv.org/html/2607.10044#S4.T3):*Revenue*is total advertiser revenue for the arm;*Coverage*is the fraction of queries for which a sponsored ad is shown;*Impression*is Impression Yield \(IY\), the expected value per impression opportunity;*DefectML\-FBS*aggregates the Defect\-ML classifier and Feedback\-Based\-Suppression \(FBS\) guardrails on the Ads Selection\.

## Appendix HMARISA\-Int Feasibility Measurements

MARISA\-Int is included only as a feasibility check \(not a primary performance baseline\)\. At billion\-key scale it is too slow for a full 13K sweep, so we report a fixed 10\-query subsample \(seed 42\) and omitK≥500K\\geq 500after observing a singleB​W=1000BW\{=\}1000query at10,99710\{,\}997s \(≈\\approx3\.05 h\)\.

Table[7](https://arxiv.org/html/2607.10044#A8.T7)reports mean/p50/p90/p95 on the subsample\. p99 is omitted atN=10N\{=\}10\. The large mean–median gap reflects high per\-query variation in traversal depth and root branching\.

Table 7:MARISA\-Int per\-request latency on a fixed 10\-query random subsample of the 13,000\-request workload \(seed 42, identical query indices acrossB​WBW\)\.Latency grows roughly linearly withB​WBW, consistent with MARISA\-Int’s serial per\-beam execution model, and confirms it is unsuitable as a main comparator at this operating point \(hence MARISA\-Opt in Section[2\.1](https://arxiv.org/html/2607.10044#S2.SS1)\)\.

## Appendix IPPT Baseline Construction

For each decoding depthd∈\[0,T\)d\\in\[0,T\), we collect the set of all token IDs that appear at positionddin any keyword of the constraint library and store them as a sorted, deduplicated array on the GPU\. A flat\(valid\_tokens, valid\_offsets, valid\_counts\)layout is used so that the depth\-ddalphabet is contiguous in device memory\. At search time, for each beam at depthdd, every top\-KKproposal token is tested againstvalid\_tokens​\[d\]\\texttt\{valid\\\_tokens\}\[d\]via a singlelower\_boundbinary search; tokens that pass are admitted into the next beam, tokens that fail are pruned\. The cooperative kernel, beam scoring, log\-prob thresholds, length normalisation, and top\-BBpruning are unchanged from FlashTrie, the only difference is the constraint structure\.

PPT admits a tokenxt,ix\_\{t,i\}at depthddiffxt,ix\_\{t,i\}appears at positionddin some keyword; it does not require that the prefixx0:dx\_\{0:d\}extended byxt,ix\_\{t,i\}corresponds to a valid path in the constraint library\. The candidate set produced by PPT is therefore a strict superset of the trie\-valid set, and PPT performs strictly less work per step than FlashTrie\.

## Appendix JFull Ablations: Binary\-Search and Linear\-Search Variants

This appendix gives the full per\-B​WBWdata for the two ablations summarised in Table[2](https://arxiv.org/html/2607.10044#S4.T2)\(PPT and the per\-key linear probe\), and then reports two additional, fully sequential GPU variants that were measured but failed to scale to the production benchmark\. Together with the main\-text PPT ablation \(Section[4\.6](https://arxiv.org/html/2607.10044#S4.SS6)\), these results decompose the FlashTrie speedup into three additive contributions: GPU residency and kernel fusion \(PPT vs MARISA\-Opt\), the LOUDS\-trie structure for path validity \(FlashTrie vs PPT\), and binary search at the inner per\-key loop \(FlashTrie vs the per\-key linear probe\)\.

### Per\-B​WBWSweep: PPT and Per\-Key Linear Probe

Table[8](https://arxiv.org/html/2607.10044#A10.T8)gives the complete per\-B​WBWsweep for the two ablations summarised in Table[2](https://arxiv.org/html/2607.10044#S4.T2): all ten beam widthsB​W∈\{100,200,…,1000\}BW\\in\\\{100,200,\\dots,1000\\\}on the same 13K\-request workload, with Mean and p95 latency in milliseconds and the speedup column equal to ablation p95 divided by FlashTrie p95 at the sameB​WBW\.

Table 8:Full per\-B​WBWablation sweep \(companion to Table[2](https://arxiv.org/html/2607.10044#S4.T2)\)\.Lower is better\.
### Why the Inner Search Algorithm Matters: Linear\-Search Variants

The PPT ablation attributes the remaining gap between FlashTrie and a GPU baseline to the LOUDS\-trie data structure\. Within the trie kernel itself, however, the per\-parent child lookup is performed by a binary search over the parent’s sorted child labels\. To isolate the contribution of this algorithmic choice, as distinct from GPU residency, kernel fusion, and the LOUDS layout, we replaced the binary\-search inner loop with three successively weaker variants and measured each in turn\. We report this ablation primarily as evidence that the seemingly small step from binary to linear search at the inner loop is responsible for a substantial fraction of FlashTrie’s wall\-clock advantage, and as a cautionary data point for naive GPU ports of the textbook MARISA traversal\.

All three variants share FlashTrie’s cooperative kernel, beam scoring, log\-prob thresholds, and length normalisation\. They differ only in how each parent beam locates the children whose labels appear in the top\-KKproposal set:

- •Two\-pointer merge \(vanilla MARISA port\)\.Faithful single\-thread port of the CPU MARISA child\-extension loop: one CTA per parent, one thread per CTA walks the parent’s sorted children and the sorted top\-KKin lockstep,O​\(degree\+K\)O\(\\text\{degree\}\+K\)per parent\. The remaining 511 threads of the CTA idle\. Across\-parent parallelism is preserved \(one CTA per parent, many CTAs per SM\), matching the CPU pattern of multi\-threading across parents\.
- •Per\-child parallel linear scan\.Block\-parallel baseline that distributes children across the CTA’s threads: threadtthandles childrent,t\+blockDim,…t,\\ t\+\\text\{blockDim\},\\ldotsand for each assigned child performs anO​\(K\)O\(K\)linear scan through the top\-KKarray to test membership\. Per\-parent workO​\(degree⋅K/blockDim\)O\(\\text\{degree\}\\cdot K/\\text\{blockDim\}\)in wall\-clock, but every thread re\-issues trie metadata loads \(select0,rank1, label decode\) for its assigned children with no sharing\.
- •Per\-key linear probe \(FlashTrie inner loop with binary to linear\)\.Identical parallelisation to FlashTrie – one thread per top\-KKkey, stridingblockDim– but each thread linearly scans the parent’s sorted children for its assigned key \(with sorted\-array early exit\) instead of binary searching\. To make this competitive at all, the CTA cooperatively decodes each parent’s child labels once into a shared\-memory cache, after which all 512 threads scan the cache\. Per\-parent workO​\(degree\+K⋅degree/blockDim\)O\(\\text\{degree\}\+K\\cdot\\text\{degree\}/\\text\{blockDim\}\)in arithmetic, with all global\-memory trie loads amortised across the top\-K loop\.

Both fully sequential variants fail to scale to the full benchmark at production beam widths\. The two\-pointer merge, although algorithmically the lowest in arithmetic complexity per parent \(O​\(degree\+K\)O\(\\text\{degree\}\+K\)\), uses only1/5121/512of each CTA’s threads; the resulting underutilisation makes the warm\-up pass alone \(20 iterations on a single batch\) take several minutes atK=100K=100, with a full 13K\-query pass projected at well over an order of magnitude longer than even the per\-child parallel variant\. The per\-child parallel scan is similarly impractical: although every thread is busy, the redundant per\-thread re\-decoding of children labels \(each thread re\-issuesselect0,rank1, and label lookups for its slice\) and theO​\(K\)O\(K\)inner scan per child push per\-batch latency into the multi\-second range\. A single 13K\-query pass atB​W=100BW=100failed to complete within seven hours, ruling out the variant for our sweep\.

Only the per\-key linear probe with cooperative shared\-memory caching admits a meaningful comparison: it shares FlashTrie’s parallelisation scheme exactly, differing only in whether the inner per\-key search over the cached children is binary or linear\. Even so, atB​W=100BW\{=\}100the per\-key linear probe has a mean per\-request latency of 110\.2 ms against FlashTrie’s 0\.55 ms \(∼200×\{\\sim\}200\\timesslower; full per\-B​WBWsweep in Table[8](https://arxiv.org/html/2607.10044#A10.T8)\), indicating that binary search at the inner loop contributes a two\-orders\-of\-magnitude factor on top of the GPU residency, kernel fusion, and shared\-memory caching that all three variants share\.

The first two variants illustrate a common failure mode of porting CPU succinct\-trie code to the GPU\. The CPU MARISA loop is fast because \(i\) modern x86 cores are individually fast and \(ii\) MARISA\-Opt multi\-threads across parents to exploit thread\-level parallelism\. Both properties evaporate on the GPU: an SM’s individual thread is far weaker than a CPU core, and the across\-parents axis alone is not wide enough to amortise the trie metadata cost when each parent’s inner loop is sequential\. The per\-child parallel variant uses all threads but pays for it with redundant global\-memory traffic that swamps the arithmetic savings\. Only by holding the parallelisation scheme fixed and varying the inner\-loop algorithm, the per\-key linear probe versus FlashTrie’s binary search, can the contribution of the search algorithm itself be measured cleanly, and that contribution is large\.

## Appendix KNQ \+ GENRE Workload Construction and Realistic\-Threshold Analysis

The three components introduced in Section[4\.7](https://arxiv.org/html/2607.10044#S4.SS7)are pinned for exact reproduction: NQ\-Open validation \(3,600 questions, free\-form input matching serving shape\), the public GENRE\-KILT checkpoint \(facebook/genre\-kilt, BART\-large with a∼50\{\\sim\}50k token vocabulary, ungated\), and the deduplicated KILT Wikipedia titles \(∼6\{\\sim\}6M tokenized titles, BART\-tokenized so the model’s emitted token IDs and the trie’s edge labels share one ID space, avoiding train/serve vocabulary mismatch\)\.

GENRE is autoregressive: a vanilla top\-KKbeam roll\-out produces a proposal grid where each position’s top\-KKis conditioned on the previous\-step argmax, so the grid is internally coherent and the trie does little admissibility filtering\. To simulate the cross\-position incoherence of production NAR grids, we use a two\-stage pipeline\.

*Stage 1: AR top\-KKextraction\.*For each NQ\-Open query we run GENRE overT=16T\{=\}16decoding positions, emitting at each positionttthe top\-K=600K\{=\}600tokens with their log\-softmax scores\. The result is a per\-query\(T,K\)\(T,K\)grid of token IDs and log\-probs, the same data shape FlashTrie and MARISA\-Opt consume on internal traffic\.

*Stage 2: head \+ noise\-tail replacement\.*For a chosen head sizer∈\{0,9,18,37,75,150,300,600\}r\\in\\\{0,9,18,37,75,150,300,600\\\}we keep the top\-rrAR\-emitted tokens at each position \(the “coherent head”\) and replace ranks\[r,K\)\[r,K\)with tokens sampled from the position\-stratified marginal distribution of the trie corpus \(the empirical frequency of each token at positionttacross all tokenized Wikipedia titles\)\. Noise tokens inherit the AR model’s tail log\-probs at the occupied ranks so the resulting grid is indistinguishable from a real model emission to a downstream beam scorer\. The knobrrinterpolates betweenr=0r\{=\}0\(pure NAR\-style proxy, maximum trie work\) andr=Kr\{=\}K\(pure AR, minimum trie work\)\.

Alternative noise designs \(uniform\-vocab, uniform\-position alphabet, and history\-conditioned sampling\) were rejected because they either under\-stress the trie or reintroduce AR coherence\. Position\-stratified marginal sampling preserves per\-position plausibility while remaining cross\-position independent\. We keep AR tail log\-probs to preserve realistic score decay underθtok\\theta\_\{\\mathrm\{tok\}\},θsent\\theta\_\{\\mathrm\{sent\}\}, and length normalization\.

All experiments use the same Docker container as Section[3](https://arxiv.org/html/2607.10044#S3)\(1×\\timesA100 80 GB, 8 CPU cores pinned vianumactl \-\-cpunodebind=0 \-\-membind=0 taskset \-c 0\-7, 60 GB host RAM\)\. The constraint trie is built once over the tokenized KILT corpus and reused across allrr\. For therr\-sweep, beam\-score thresholds are set permissively \(θtok=−20\\theta\_\{\\mathrm\{tok\}\}\{=\}\{\-\}20,θsent=−200\\theta\_\{\\mathrm\{sent\}\}\{=\}\{\-\}200, length\-norm exponent5\.05\.0\) so that the trie’s admissibility check, not the scorer, drives all pruning\. Results are reported atB​W=600BW\{=\}600, which sits inside the beam\-width range of the internal sweep \(Figure[3](https://arxiv.org/html/2607.10044#S4.F3)\) and has the densestrr\-grid in the released artifacts\. The latency measurement protocol \(Section[3](https://arxiv.org/html/2607.10044#S3)\) is reused verbatim\.

Table 9:Mean trie\-search latency \(ms\) atB​W=600BW\{=\}600on NQ\+GENRE vs\. head sizerr\(r=0r\{=\}0: NAR\-style proxy, max trie work;r=600r\{=\}600: GENRE verbatim\)\. Speedup = CPU / GPU mean\.Lower latency, higher speedup are better\.### GPU Overhead Floor

The GPU curve flattening in Table[9](https://arxiv.org/html/2607.10044#A11.T9)is explained by a structural overhead floor \(kernel launch, grid barriers, fixed top\-BBselection passes, and fixed\-size H2D/D2H I/O\)\. CPU has less fixed overhead, so its latency keeps scaling with reduced trie work while GPU saturates once work drops below this floor\.

### Realistic\-Threshold Operating Point

Table[9](https://arxiv.org/html/2607.10044#A11.T9)uses permissive thresholds to stress trie work\. Here we evaluate realistic production thresholds \(Section[4\.8](https://arxiv.org/html/2607.10044#S4.SS8)\), where scorer\-side pruning is stronger under the same 30 ms SLA\.

We keep all other NQ\+GENRE settings fixed \(3,600 queries,B​W=600BW\{=\}600,T=16T\{=\}16, length norm 5\.0, identical warmup\) and use the deployed threshold pair\. Table[10](https://arxiv.org/html/2607.10044#A11.T10)reports mean latency and CPU/GPU speedup across therrsweep\.

At realistic thresholds and worst\-caser=0r\{=\}0, MARISA\-Opt drops from 41\.19 ms \(permissive\) to 8\.95 ms mean and FlashTrie from 5\.66 ms to 1\.54 ms, showing scorer\-side pruning reduces absolute work for both systems\. Relative advantage remains multi\-×\\times: 5\.8×\\timesatr=0r\{=\}0, typically 7\.4–8\.3×\\timesforr∈\{9,37,150,300\}r\\in\\\{9,37,150,300\\\}\(Table[10](https://arxiv.org/html/2607.10044#A11.T10)\)\.

Across therrsweep, CPU latency still tracks workload coherence, while GPU latency stays in a narrow band because fixed cooperative\-kernel overhead dominates once useful trie work is small\. This preserves FlashTrie’s core value at production thresholds: lower and more stable latency under workload variance\.

Table 10:Per\-request mean trie\-search latency \(ms\) atK=600K\{=\}600under realistic production thresholds on the NQ \+ GENRE public workload, across the head\-size sweepr∈\{0,9,18,37,75,150,300,600\}r\\in\\\{0,9,18,37,75,150,300,600\\\}\. MARISA\-Opt is the CPU baseline \(8 worker threads, NUMA\-local pinning\); FlashTrie is the GPU implementation on a single A100 80 GB\. TheSpeedupcolumn reports CPU mean / GPU mean per row\. Compared to the permissive\-threshold sweep in Table[9](https://arxiv.org/html/2607.10044#A11.T9), score\-driven pruning ahead of the trie reduces CPU mean latency by≈5×\\approx 5\\timesand GPU mean by≈2–4×\\approx 2\\text\{\-\-\}4\\times, while the relative CPU\-to\-GPU mean speedup remains in the same55–13×13\\timesband\.Lower is better for latency; higher is better for speedup\.