Repository navigation
Add AI.CLASSIFY with document KV reuse across query operations - #178
Merged
Merged
Conversation
A new cell scores each classification label as the sum of its tokens' log probabilities under Qwen3 32B fp8 on vLLM 0.26.0. It sends one request per shared label prefix and reads the next tokens' log probabilities with logprob_token_ids, which keeps prefix caching; prompt log probabilities would bypass the cache and recompute each document once per label. Each document's first request fills the cache before the others run. The cell labels sf=1.0 in shards, copies labels by document content to sf=0.5 and sf=0.1, and extends the published collections, reusing every filter and join label set. activate_reused_collection gains new_specs for predicates outside PREDICATES. quail-b moves to the commit that defines the classification queries.
Quail's runtime has no AI.CLASSIFY yet, so its QUAIL-B runner leaves classification queries out of default runs and reports an explicitly requested one as skipped, instead of failing to read its plan.
The classification queries are now in quail-b's default list. Quail's runner still skips them, and the tests expect them in QUERY_ORDER but not among the queries Quail plans. The labeling cell's collection check selects queries by their classifications, since none is pending any more.
…n-atl3of # Conflicts: # tests/test_quailb.py
fsdatalab/quail-bench#31 was squash-merged as b47d212, with the same tree as the branch commit 3738036 pinned before.
A classification is a column-producing model call, planned and run like AI.SCORE: each document is a join anchor whose prompt head and document stay in KV, the classification question and category list are written once after it, and one suffix per label-trie node reads the next token's log probabilities over the whole vocabulary. Each label's score is the sum over its tokens; a tie goes to the earlier label. A label filter keeps documents whose label is accepted. Quail keeps the full output head only when it is tied to the embedding, so Qwen3 4B classifies and Qwen3 32B refuses. Joins, AI.IF filters beside a classification, diffusion models, and the request backends refuse for now. The QUAIL-B runner reads ai_classify plans and label filters, and runs IMDB-11, IMDB-14, BIO-5, FEV-11, and AGENT-4.
The labeling cell's judge takes a model, and quail_check labels the documents of a QUAIL-B run with stock vLLM on Qwen3 4B, then reports which labels differ from Quail's and by how much vLLM's scores separate them.
The cell scores the documents where a Quail run and stock vLLM chose different labels, and a sample where they agreed, four ways: Quail's frame path, Quail with the question in the anchor prefix, Quail with float32 logits, and vLLM's full-sequence prompt log probabilities.
shreyashankar
changed the title
Add the label_chains scoring rule: one request per label, holding the answer cue and all but the label's last token, with every row's log probabilities read. The trie_nodes rule stays selectable, so the two can be compared on the same query. - The executor's pack_chunk takes read_all_rows and records rows_per_answer; AsyncLabelLogprobs pads each suffix's rows into one fixed-size record, so JoinAdmission.report needs no change. - ClassifySpec names its scoring rule; EngineConfig.label_scoring forces one for a session, and the planner records it in the plan settings. The planner costs each rule by its own suffix lengths. - The classify planner is a plain loop over steps with a small table record, without the closure it had. - The bench runner and the parallel Modal runner take --label-scoring.
When every label is one token the readout holds one row per chain and returns (labels, targets); reshape it before chain scoring. IMDB-11's sentiment labels hit this on the first label_chains run.
score_in_batches replaced the extension metrics with the output name, aliases, and input rows, so label_tokens never reached the report. Sum it over batches for AiClassify nodes.
The prefix_sharing rule now fires for a classification as it does for a join: its documents' KV is written anyway, so any whole shared page pays. The classifier builds the prefix tree over its document prefixes and hands it to run_join, which packs a child from its shared length on. A classification that shares prefixes scores every row in one call, so parents and children meet in the same tree. Borrowed tokens are reported per node as borrowed_prefix_tokens.
Add the trie_paths scoring rule, now the default: one chain per deepest proper prefix of the label trie, every row read and every label token returned. The answer cue's row scores all one-token labels at once, so a one-token label set costs one request per document, and labels sharing leading tokens share the rows after them. The rules' requests and scoring are one LabelRequests record per rule.
shreyashankar
marked this pull request as draft
September 28, 2026 19:26
The compile and touch passes now run a classification-shaped join chunk: the question as each document's frame, then 26 suffixes of one to six tokens. WARMUP_VERSION moves to 5 so the marker forces one compile pass. The bench records the file count under each kernel cache directory before and after every Quail query, so a kernel compiled inside a query's time shows as growth.
…bench The classifier now turns its forward chunks' CUDA events into gpu_s and chunks on the node's metrics when the session asks for GPU timing, as filters and joins do. The bench CLIs take --gpu-timing to ask for it.
A classify node now carries later stages: classifications of the same documents, each gated by the accepted labels of the label filter before it. The executor runs every stage in one join call, so a document goes on to the next stage while its head and document are still in KV, packing only that stage's frame and label suffixes; the join admission takes an advance predicate that reads the gate from the stage's label. The stages share one readout over every stage's label tokens. Later labels come back as extra columns of the scores table, and each stage gets its own classify answer table. The planner folds a classification into the node before it when both run on the same alias with the same prompt head and at most one label filter, on the earlier label, lies between them. Its estimate streams a later stage's frame and suffixes against a resident document. The score node's inputs now carry the query's gpu_timing choice, so classify nodes report gpu_s and chunks like filters and joins.
shreyashankar
changed the title
A Stage is what one model operator asks of each document at one step: its frame, its request suffixes, the readout for its answers, which suffixes each document sends, and the decision that lets a document go on. run_stages drives a list of stages over documents with one admission, so a document's KV is computed once and stays resident while it moves through every operator's stage. Each stage has its own readout; a frame the stage before already wrote is not written again; a survivor limit ends admission; an unpaged mode runs a lone single-suffix stage without arena pages. run_join is now a wrapper that builds one Stage per join stage, so joins, AI.SCORE, and classify chains run on the scheduler unchanged.
filter_stages builds one Stage per question: the preamble the questions share is the frame, written once at the first stage, and each stage's suffix is its question past the preamble. A single-suffix stage packs the frame and its suffix as one entry written straight into the document's pages, and under tree attention a borrowing document reads its parent's pages stacked with its siblings, as the filter loop did. run_filter delegates to run_stages: the unpaged path for a lone stage, the survivor limit, retention through on_settled, and the per-chunk finished-document callback through on_chunk.
A filter chain that feeds a join now runs as leading stages of the join's own run: every document of the chain is an anchor, its questions are the first stages, and a survivor goes on to the join's stages with its KV resident. The join's per-document partner lists are asked when the document reaches the stage, so a per-batch transform between the chain and the join runs on each survivor there and can drop it. The join's anchors are the documents that reached it, in the chain's order. The filter stream and the filter admission are gone: run_stages with JoinAdmission is the scheduler for filters, joins, AI.SCORE, and classifications, on the worker's child-join path as well. Tests that built the stream now build stages.
A family whose queries Quail cannot run yet, such as the classify and join queries, left run_suite with no query ids and quail-b refused the run. The suite now returns a record with no queries and the skipped ids, in the shape the parallel runner merges. Add experiments/cells/profile_query.py: one QUAIL-B query under cProfile on Modal, to read the host time between forward passes.
The cell shares the parallel runner's Modal app, whose entrypoint is already called main.
- The bench counted kernel-cache files inside the timed window, and walking the caches on the volume took seconds. The count now runs before the query's clock starts and counts top-level entries only. - The readout head is kept in bf16: the fp32 head made every full-vocabulary matmul an fp32 GEMM, 2.8 of BIO-5's 3.5 GPU seconds. Logits are computed in bf16, as vLLM computes them, and normalized in float32. - A chunk's token parts are joined on the host and copied into pinned memory once, instead of one copy per suffix: 41,636 copies on BIO-5.
The Quail physical planner chose each table's AI.IF filter order, the joins' stage order and anchors, and which KV stays resident while it built the graph, so EXPLAIN could not name those decisions and no extension could see them. They are rules now. Logical rules, in order: projection_pushdown, filter_pushdown, classify_placement, filter_order, and join_order. filter_order records each SemanticFilter's execution order on the node; join_order runs the left-deep search once per plan and records each SemanticJoin's execution position and anchor (exec_idx, exec_anchor), leaving the written anchor as the constraint it was. The search result is memoized on the planning context, so classify_placement's candidate pricing and the registered rule share one search. Physical rules, in order: kv_retention, label_scoring, prefix_sharing, and tree_attention. kv_retention computes the retention schedule from the decided stages, sets keep_kv, arena_writes, and keep_anchor_kv, and puts the schedule in the plan's retention setting through the context; it fires once, so a plan refined on exact tokens keeps the schedule it was made with, as before. plan_quail now only translates the decided logical plan into physical nodes, with its refusals. The filter-order cost code moved to quail.planner.filters and the shared statistics to quail.planner.statistics. Physical plans are unchanged for every QUAIL-B query at sf 0.5 on qwen3-4b-fp8 and diffusion-gemma-26b-a4b-fp8, first and refined, except for the rule remarks.
A plan carried remarks: a line for each rule that changed it, and advice such as a forced join anchor that priced higher than a free one. EXPLAIN printed them under its settings. Every plan named projection_pushdown, so the lines said little, and Spark keeps rule changes out of a plan too (spark.sql.planChangeLog.level is an opt-in log). PhysicalPlan, PlanningContext, and LogicalPlanningContext lose their remarks, plan_query its remarks argument, and join_order no longer runs a second search with forced anchors released, which only priced that advice. A run report's remarks keep the session's notes. Physical plans for every QUAIL-B query at sf 0.5 on Qwen3 4B and DiffusionGemma are identical to 3b0946a. Tests that read remarks now check what each rule set on the plan.
A filter on an AI.CLASSIFY label makes no model call, so the logical plan now holds it as a Filter node with an InList condition on the column the SemanticClassify adds, instead of a LabelIn predicate inside a SemanticFilter. SemanticFilter holds model calls only. LabelIn is removed. The planner reads label filters through Operators.label_filters (LabelFilter: the call, values, selectivity, and written position after the table's AI.IF predicates), so physical node ids are unchanged. Logical EXPLAIN prints the condition, such as "Filter: m.organ_class IN [...]", where it printed the classification's prompt before. Physical plans of the 45 quail-b queries, planned twice each, and of 12 synthetic classification cases are identical before and after.
Moves the quail-b pin to fsdatalab/quail-bench#35, which adds IMDB-16 and its five guideline checks as labels-pending filter predicates. experiments/cells/imdb16_labels.py labels each check over the sf=1.0 reviews on its own H100 with Qwen3 32B through Quail, copies the labels to sf=0.5 and sf=0.1 by review content, and builds collections that reuse every other published label set. The IMDB filter group now has eight predicates, so its label parts hold fewer rows; the labeling test derives the part size from the group instead of assuming three filters.
run_stages took its progress unit as a string that also chose what to count: any unit but "anchors" counted answers. The unit is now only the label, "documents" by default, and count_answers chooses the count; AI.SCORE over pairs still counts scores.
Moves the quail-b pin back to b82b14f and removes experiments/cells/imdb16_labels.py. IMDB-16 (fsdatalab/quail-bench#35) is dropped: its five checks repeat one guidelines text only to create a shared prefix, which no analyst would write. This reverts 4508215.
The decoding baseline feeds every generated token but the last back through the model, so those tokens are fresh work. They were recorded as generated_tokens but left out of fresh_tokens, which made the baseline's KV regret too low.
A family split into groups (one BIOdex query per container) saves each group in its own folder, such as biodex-bio-5. The merged run.json named the family folder instead, so rescoring the merged run did not find the query files.
The measured time is wall-clock time and is over 1 ms on a busy machine.
… starts fsdatalab/quail-bench at f3ad796 counts the partner label once per anchor and the partner documents of one anchor as a prefix trie.
Every pair of one anchor has the same partner label directly after the frame, so the label is now part of the frame. Each pair writes only its partner document and the answer cue. The token sequence is the same, so the answers do not change. The joined-row classification, the join cost, and the speed of light estimate count the label the same way.
Describe what Quail does instead of giving commands, keep the details a user or an implementer needs, and bring the logical plan and planning pages up to date: the label Filter with InList, the five logical rules, the planner steps, and the physical rule order.
…nto claude/join-label-in-anchor-kv
shreyashankar
added this pull request to stack #183
October 3, 2026 00:11
- Docstrings follow Google style and describe current behavior; history, run results, and PR references are removed. - Comments that repeat the code are removed. - Comments and docstrings that contradicted the code are corrected. - Classification text uses "label" and the rule names letters, trie_tree, and trie_decode. - A request-backend test name matches what it checks: one decode request per document.
…nto claude/join-label-in-anchor-kv
- Architecture follows a query: logical plans, optimizer, physical plans, cost model, execution, KV. - Optimizer lists all nine rules in run order, with one section per cost-based rule. It replaces the rule sections of Planning and Logical plans. - Cost model describes the roofline model at a high level and links to the blog post for the filter derivation. - KV replaces Prefix caching and covers paged KV, KV rewind, prefix sharing, and KV retention. - Extending > Rules shows a logical and a physical rule, registration, the planning contexts, where to put a rule, and how to test it. - Extending > Overview describes load_extension.
…nto claude/join-label-in-anchor-kv # Conflicts: # docs/content/docs/architecture/planning.mdx
shreyashankar
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Quail could filter and join documents with
AI.IF, but it could not assign classification labels within a query. Queries with several model operations could also process the same document repeatedly after its KV was released.Solution
Add
AI.CLASSIFYto SQL and the Python builder so queries can assign labels, filter on labels, and classify rows produced by a join. The PR also adds 12 classification queries to quail-b to evaluate the feature.Refactor execution to reuse document KV across consecutive model operations. KV is the model's saved attention keys and values for tokens it has already processed. Keeping document KV in GPU memory avoids processing the same document again for each operation.
For example, imagine a table of movie reviews. A query can filter reviews, assign sentiment labels, and join selected reviews with movie aspects while reusing the review KV.
flowchart LR scan["Read reviews"] --> filter["Filter reviews with AI.IF"] filter -->|"Reuse review KV"| classify["Assign sentiment with AI.CLASSIFY"] classify --> labels["Filter on sentiment labels"] labels -->|"Reuse review KV"| join["Join reviews with aspects using AI.IF"] join --> result["Return selected columns"]Experimental results
Across the 12 new classification queries, total query time was 2,781 seconds on Quail, compared with 5,676 seconds on stock vLLM. Both used Qwen3 4B FP8 at scale factor 0.5 on one H100. Stock vLLM submitted one operator at a time with prefix caching. Query times exclude model startup. At the configured price of $3.9492 per GPU hour, total GPU cost was $3.05 for Quail and $6.23 for stock vLLM.
No stock vLLM comparison is recorded for DiffusionGemma.
Agreement with the Qwen3 32B reference labels was lower on AGENT-3 at 0.833 for Quail, compared with 0.886 for stock vLLM. On BIO-5 and BIO-6, agreement was 0.720, compared with 0.812. The vLLM calculation includes only rows with a recognized label. The cause remains unresolved and needs comparison across scoring methods.
Measurement sources
Results are on the
quail-resultsvolume. The measured commits precede the latest change to join execution. That change has not been measured.f15c5ca2, rescored with the corrected token minimum./results/benchmarks/quailb/20261002T061804Z-a6ccc1fe-minfix/quail6deb4fc, rescored with decoded tokens included in fresh tokens./results/benchmarks/quailb/20261001T010217Z-1e7f7ba8-minfix/stock_vllmf3ad796. The current dependency pin is4298e7a, the merge commit for that benchmark change.Checks
954d5f39. Ruff, long string checks, Vulture, and pytest pass. The local suite has 201 tests passed and 9 skipped.Closes #105