Optimizer

The optimizer chooses how Quail will execute a query. It starts with a logical plan, which describes the requested operations, and produces a physical plan, which specifies the executable operators and their settings.

Quail makes planning decisions through optimizer rules and while building the physical plan. Some rules remove unnecessary work, such as loading columns the query does not use. Other rules compare alternatives using the cost estimates and choose the one with the lowest estimated GPU time.

This page describes the built-in rules. See Rules to write and register an extension rule.

Planning process

Query.plan() performs these steps:

1. Get document lengths in tokens. If a column is not yet tokenized, estimate its lengths using the tokens per byte of the first 1,024 documents.
2. Run the logical rules in order. Repeat the sequence until a pass makes no change, or four passes have run.
3. Build the physical graph with build_physical_plan. Set the token budget and KV capacity from the model and device specifications, and choose classification placement, filter order, and join order.
4. Run the physical rules once, in order.

Tokenization can continue in the background while planning runs. Before execution, refine_plan applies the physical rules with exact token counts and sequences. It preserves the choices already made while building the plan and by rules such as label_scoring and kv_retention.

Cost estimates

The cost model estimates GPU time from the computation and memory transfers a plan requires. It assumes peak GPU compute rates and memory bandwidth and excludes CPU work and kernel launch time. Use explain(analyze=True) to compare an estimate with measured execution.

The model computes key and value vectors when it processes a token. Quail saves the vectors as KV for reuse in later requests. The cost model counts:

For example, imagine two questions about the same document. Keeping the document's KV lets the second question reuse it. The model processes the new question tokens and reads the saved KV without processing the document again.

The roofline model

Model dimensions convert the counts into floating point operations (FLOPs) and bytes transferred. GPU specifications give operations per second and bytes transferred per second. Quail uses the larger of the computation and memory transfer times:

time_seconds = max(FLOPs / peak_FLOPs_per_second,
                   bytes_transferred / bytes_per_second)

The calculation is called the roofline model. Quail totals the work across layers for each part of the model calculation, applies the formula, and adds the resulting times.

Our blog post on filter cost estimates explains the equations, GPU background, KV reuse, and filter ordering. Detailed explanations of join and classification costs are TBD.

Speed of light estimate

quail.speed_of_light_estimate estimates query time under ideal execution conditions for comparison with measured runs. It assumes:

• One GPU running at peak compute rates and memory bandwidth.
• Unlimited KV capacity. By default, each distinct document prefix is computed once and reused.
• No time spent on CPU work, scheduling, or kernel launches.
• Rows that pass filters and joins are determined by answers supplied by the caller, such as saved reference answers and classification labels.

For classification, this estimate counts the question and category text but excludes the extra steps used to select a label. The optimizer's classification estimates include the selected label scoring method.

Report the supplied answers and KV assumptions with the estimate. If a measured run produces different answers, it can evaluate different rows and perform different amounts of work.

Built-in rules

Built-in rules run in the order below. Extension rules run after the built-in rules of the same kind, in registration order.

The implementations are in logical_rules.py and physical_rules.py.

Building the physical plan also makes three cost-based choices, described below.

Each operator supplies its own prompt preparation, execution constraints, and costs. Score operators support both generative models and reranker models. Score joins preserve the document order in the model prompt. The executor reuses the first document's KV across its partners, so that document is the anchor.

Call explain(verbose=True) to inspect physical settings, including share_prefixes and attention, and to list the rules with the cost-based ones marked.

Physical plan choices

build_physical_plan and the Quail backend make these choices while building the physical plan. They compare alternatives using estimated query work.

Classification placement

The Quail backend decides when to classify a table's documents. A classification that reads one table's document can run before or after joins involving that table. This differs from classifying a joined pair, where the classification reads both documents.

• Before the joins, classification evaluates every document that passes the table's AI.IF filters. A condition on the label can reduce the number of documents entering the joins.
• After the joins, classification evaluates only documents that survive those joins.

For example, imagine a query that classifies reviews by sentiment, keeps negative reviews, and joins them with movie aspects. Classifying first can save join work if few reviews are negative. Joining first can save classification work if few reviews match an aspect.

The backend compares the estimated query times. It keeps the written plan unless moving classifications after the joins has a lower estimate.

Filter order

The planner chooses predicate order within a table. When a table has several model predicates, their order affects how much work the query does. Selectivity is the estimated fraction of rows that pass a predicate. A predicate that keeps fewer rows can reduce the work of later predicates.

For AI.IF, the first predicate processes each full document. With KV rewind, later predicates process only their question tokens for documents that passed. The planner accounts for that difference when comparing orders.

The planner tries each predicate first. It orders the remaining predicates by cost per rejected document and keeps the candidate with the lowest expected time. Our blog post on filter cost estimates gives the derivation for AI.IF.

For AI.SCORE, the planner prices each prompt at the expected survivor count. It accounts for prefix reuse within a score evaluation. Later comparisons of a score already computed for those rows need no model work. The physical graph builder reads the chosen order and reuses the computed score column.

Set order="as_written" to preserve predicate order.

Join order

The planner chooses the sequence and anchors of joins. The anchor is the side whose document KV stays in GPU memory while its pairs are evaluated. The partner is added after the anchor for each pair.

A long anchor document can be processed once and reused for many partners. A long partner document must be processed for each anchor it pairs with. The planner considers this work along with filtering and KV retention.

The planner compares join orders and anchors by estimated time for the whole query, including filters. It considers only plans where each join adds one table to the preceding result.

An explicitly supplied anchor is preserved. With order="as_written", the planner preserves join order and chooses only the unspecified anchors.

Physical rules

Physical rules choose execution settings for the graph. They use the model and device specifications together with the execution limits below.

Execution limits

build_physical_plan sets two execution limits:

• The token budget is the maximum number of new tokens in one forward pass. Available GPU memory and kernel index limits determine this budget.
• KV capacity is the number of KV pages that fit in GPU memory after reserving memory for model weights and execution.

See Execution for how these limits affect batching.

KV retention

The kv_retention rule identifies document KV that later operators can reuse. A filter can retain the KV of passing documents for a later join anchored on that table. A join can retain anchor KV for a later join with the same anchors.

Retained KV has a memory limit. When pages run out, the executor removes the KV with the lowest estimated future savings per page. An operator that later needs the removed KV must compute it again.

See KV retention for the priority calculation and execution behavior.

Label scoring

The label_scoring rule compares supported classification methods and chooses the one with the lowest estimated GPU time. It considers the model, label tokens, requested output, and whether earlier operators already provide document KV.

Decision 2.0 Kai uses its own scoring method, decision_choice, so the rule does not compare alternative methods for that model. See Label scoring for how the methods select labels and their requirements.

Prefix sharing

The prefix_sharing rule enables prefix sharing when the estimated compute savings exceed the additional KV write time:

saved_seconds = shared_tokens * 2 * model.params / device.peak_flops
write_seconds = total_tokens * model.kappa / device.hbm_bw

If the operator already writes KV, the additional write time is zero. Joins and classifications already write their document KV.

The comparison excludes sorting, scheduling, attention, and recomputation when shared KV is unavailable. Prefix matching requires token sequences, so the rule can make its final decision after tokenization finishes.

Attention path

The tree_attention rule chooses between the unified and tree attention paths described in Execution. It compares their estimated time for requests that read the same prefix. For r requests, each with q new tokens and a shared prefix of s tokens, it compares:

unified_seconds = r * attention_seconds(q, s)
tree_seconds = attention_seconds(r * q, s) + merge_seconds

attention_seconds(q, s) estimates the compute and memory time for the new tokens to attend to the saved prefix, summed across layers. merge_seconds estimates the memory time to read two partial attention outputs and write the combined output.

• For joins, the inputs are estimated partners per anchor, tokens per partner, and mean anchor length.
• For filters, the rule compares paths for each group of documents sharing a parent prefix. It chooses the path preferred by the most reused tokens. Ties choose unified attention.
• Tree attention supports BF16 and FP8 weights. It requires at least two billion parameters, excluding embeddings and output layers, and no extra answer positions. Other configurations use unified attention.

The current comparison uses the FP8 compute rate for attention, although the kernel runs in BF16. It therefore underestimates attention compute time on an H100. It also excludes kernel launch time. These limits apply to the path comparison, rather than the main roofline attention estimate.

The implementation is choose_attention_path.