Sharding: up to 256TB in one index

Post-training open-weight models for large-scale code search

September 04, 2026Edward Yu (Applied Compute), Jacob Phillips (Applied Compute)
guest

Today’s frontier models are exceptional at code search. They have been extensively trained to use nothing more than traditional human tools: the filesystem, ls, and grep. This allows many small-scale code search use cases to avoid the complexity of building a dedicated search index.

However, this simple approach to code search is expensive and slow when the corpus is large. In these cases, the savings gained from eschewing search indexes can be wiped out by high inference costs and slow grep tool calls.

Applied Compute partnered with turbopuffer to show how to solve both of these problems:

  1. Train a small, specialized search model: Today’s small open-weight models struggle with search tasks out of the box, but they can be post-trained with search tools to achieve impressive search quality at a fraction of the cost and latency of most frontier models.
  2. Build a dedicated search index in turbopuffer: While using grep directly is convenient, on large corpora, turbopuffer search tools can be orders of magnitude faster.

Training setup

We create a reinforcement learning (RL) environment to train Qwen3.6-35B-A3B to use search tools to achieve different search tasks across a large multi-codebase corpus.

Data and task construction

We index ~9,000 real open-source GitHub repositories spanning many languages and project types. Each training task has a random subset of codebases assigned to it, including the target codebase(s) needed to complete its task. We test two task types, narrow and open-ended, to evaluate whether the agent can search both deeply and broadly.

Narrow task

The agent is given a corpus of up to 1,000 repositories and asked a question about one of them. The target repository is anonymized, and the question is specific enough that arriving at the correct answer requires reading code and understanding implementation, not just skimming documentation, for example:

We have a stateless database dump tool that can compress and encrypt dump
parts before storing them. What is the exact order of compression/encryption
on write and decryption/decompression on read, and how is the AES key material
and nonce derived from the configured key string?

The narrow task tests the agent’s ability to pinpoint the target repo in a very large set of codebases (a needle-in-a-haystack search) and then pull exact details out of its code.

The reward is a gated combination of:

  • Correctness: each question we generate comes with a rubric graded by an LLM-as-judge (GPT-5).
  • Citation accuracy: the model is asked to make specific citations to code blocks in its answer; we grade it for how well its claims are supported by these citations alone.
  • Efficiency: we apply a penalty in the reward term to incentivize short answers reached in few agent turns (more below).

Open-ended task

The agent is given a corpus of up to 100 repositories and asked to find examples within the corpus that implement some described code pattern (usually a design pattern or architectural idiom). Only a handful of the repositories in the corpus actually have the pattern, and the task is to find them and confirm each one by reading its code, for example:

Find as many examples of the "visitor/double-dispatch" pattern in the corpus
as you can. This is a code pattern where operations are externalized from a
class hierarchy via double dispatch; each node exposes accept(visitor) and
calls back the visitor's type-specific visit method, so new operations are
added as visitors without modifying node classes.

This task rewards the ability to efficiently sweep the whole corpus, and to search in concept space rather than code-implementation space. Matches are on an abstract shape, and it is much harder to identify specific strings that an agent might grep for, which should make this task very challenging for a bash-only agent.

The reward for the open-ended task is precision-focused. Each repo cited in the answer is assessed independently by an LLM judge that sees the cited code in context and counts it only if the citation resolves to real code and the judge agrees it implements the specified pattern. The task score is precision — the fraction of claimed repos that are correct — multiplied by the same efficiency penalty as the narrow task:

reward=precision×efficiency penaltyprecision=correctmax(claimed,K)\begin{gathered} \text{reward} = \text{precision} \times \text{efficiency penalty} \\ \text{precision} = \frac{\text{correct}}{\max(\text{claimed},\,K)} \end{gathered}

We build each open-ended task so that (at least) KK of its repos implement the pattern, and the model is told that there are KK repos to find. The expression in the denominator is chosen to encode both precision and recall: finding fewer than the KK total caps the score while padding the response with junk repos hurts it. This ensures the agent must search broadly but is only incentivized to return repositories it feels sure about.

Both tasks apply the same efficiency penalty: it rewards short answers reached in few agent turns, and is the average of a length term and a turn term:

length penalty=1min(max(LL0,0),B)Bturn penalty=f+(1f)(1min(n,N)N)efficiency penalty=length penalty+turn penalty2\begin{aligned} \text{length penalty} &= 1 - \frac{\min(\max(L - L_0,\,0),\,B)}{B} \\ \text{turn penalty} &= f + (1 - f)\left(1 - \frac{\min(n,\,N)}{N}\right) \\ \text{efficiency penalty} &= \frac{\text{length penalty} + \text{turn penalty}}2 \\ \end{aligned}

where LL is the answer length in characters and nn is the number of agent turns. The length term gives full credit up to a length L0=2048L_0 = 2048 and decays to 00 by L0+BL_0 + B (B=6144B = 6144); the turn term decays from 1 toward a floor f=0.4f = 0.4 as turns approach a budget N=20N = 20 for the narrow task and N=40N = 40 for the open-ended task. An empty answer scores 0, and one that never searched (n1n \le 1) is capped at 0.1.

Harness

The model is given a code search harness inside a read-only sandbox. The baseline harness provides a standard set of read-only exploration tools to navigate the filesystem: read_file (reads a file, optionally a specific line range), ripgrep (recursive regex search), glob (matches files by glob pattern), and list_dir (lists a directory’s contents). We then additionally provide a tool to perform hybrid search over precomputed indexes stored in turbopuffer. We use tree-sitter to chunk repositories in an AST-aware fashion, then index each chunk in two ways:

  1. a BM25 keyword index
  2. a dense embedding of each chunk, produced by the open-source Octen-Embedding-8B embedding model

Alongside the text and vectors, we store the filepath and line numbers of the chunk as metadata. At search time, we retrieve candidates with turbopuffer, rerank them with the open-source Qwen3 Reranker 4B, and use the metadata so the agent can jump straight to the source file and cite it directly.

Search tool modes

We train and evaluate our agent under two search modes:

  1. bash-only: the agent has access only to the filesystem and bash primitives.
  2. bash + turbopuffer search: in addition to bash primitives, the agent can search the precomputed turbopuffer indexes. Rather than embedding the question directly, the agent writes a hypothetical code snippet — its guess at what the answer looks like in source — which we embed with the document prompt, fuse with a BM25 leg over the same snippet, and rerank.

Training results

We blend the two task types in equal proportion and train Qwen3.6-35B-A3B with GRPO.

Training increases search quality

Regardless of task or tools, the model becomes more effective at completing the search task as it learns to use the tools available to it.

Loading chart data...

On the narrow task, the bash-only agent improves its correctness × citation support by 38% over training, while the agent with turbopuffer search improves by 57%. On the open-ended task, the bash-only agent increases precision at peak by 182% and the turbopuffer agent by 211%. Training helps either way — but the agent with access to turbopuffer search tools is learning from better evidence at every step, and it finishes 16% ahead on the narrow task and 25% ahead on the open-ended one.

Turns decrease over training

In addition to achieving better search results through training, the agent also learns to do so more efficiently, reducing turns over the rollout. This is especially pronounced in the open-ended task, where the bash-only agent ends training using 50% fewer turns than it did at the start, and the agent with turbopuffer search tools uses 70% fewer.

Loading chart data...

Search becomes more efficient with training

Diminishing turn counts show that the agent gets more efficient over training. We can see exactly how this occurs by looking at the mix of tool calls as we train:

Tool calls per rollout over training

Loading chart data...

Loading chart data...

Training pushes the two agents in opposite directions. The bash-only agent responds by grepping harder — on the narrow task its ripgrep calls more than double over training (4.0 → 8.9 per rollout), and its total tool calls actually rise, from 11 to 12. The agent with turbopuffer search goes the other way: ripgrep all but disappears (1.7 → 0.2 per rollout on the narrow task, 15.3 → 0.3 on the open-ended one) and total calls are reduced by half on the narrow task and by 70% on the open-ended one. It also learns to search better rather than more — turbopuffer search calls on the open-ended task drop from 12 per rollout to 5 even as precision climbs.

Usage of filesystem primitives like glob and list_dir decreases in both modes, but only the agent with turbopuffer search can afford to give up ripgrep: for the bash-only agent, ripgrep is the effective search index, and it finishes training sending 70% of its calls there.

Comparison to frontier models

We compare the base Qwen3.6-35B-A3B checkpoint to its trained counterpart as well as several frontier models to determine the Pareto frontier of quality versus latency and cost. All models have access to bash and turbopuffer search tools in their harness. Before RL, Qwen3.6-35B-A3B is not a very good search agent. The training is what makes the specialist.

Loading chart data...

Loading chart data...

Our fine-tuned model tops the narrow task outright in correctness × citation support. While it does not achieve frontier-level search quality in the open-ended task, it closes the gap considerably versus the base model. In both cases the specialist runs at 10–100x lower cost and 2–10x lower latency than most frontier models. GPT-5.6 Luna offers an exception, as it outperforms our fine-tuned model in cost and quality, though not latency, on the open-ended task.

Search latency scales more efficiently with precomputed indexes

One of the most compelling reasons to employ search tools over precomputed indexes, regardless of model, is how efficiently queries scale over the size of the corpus. This holds true in our training setup. Below we show the median rollout latency on the open-ended task, broken down by tool, when we scale the size of the corpus from 20 repositories to 300 repositories.

Median rollout latency by corpus size

Loading chart data...

Loading chart data...

Model completion time is similar regardless of the size of the corpus or the tools available. In contrast, a 15x increase in corpus size causes ripgrep latency to increase by 11x, while turbopuffer search latency increases by only 1.2x.

Inference cost is not total cost

A small model trained to use search tools over precomputed indexes can run search tasks at up to 100x lower token cost than most frontier models. It is important to note, however, that the full cost of the search pipeline will be determined by several factors beyond inference:

  • The size of the corpus: Larger corpora are generally more expensive to store, index, and query.
  • The write-read ratio: Precomputed indexes effectively cache, at write time, compute that would otherwise be spent at read time. When the corpus changes infrequently and sustains a high query rate, the cost of indexing is well-amortized across many fast and cheap queries. When that ratio flips, however, the cost of maintaining the index can outweigh the savings at query time.
  • The kind of search tools used: BM25 full-text search requires little upfront compute and zero inference, whereas dense vectors require inference for embeddings and increase storage demands. Techniques such as multi-vector embedding and late interaction can further increase inference and storage costs.

Conclusion

Post-training a small, open-weight model to search over a precomputed index yields impressive cost and efficiency improvements for code search over many repos. At 300 repos, this approach makes the median search nearly 3x faster than a model using only bash primitives, while decreasing the marginal cost of a search by up to 100x compared to frontier models. You can try it yourself in this demo.

While our test harness focused on code search, we expect these results to generalize to any large-scale corpora.

turbopuffer

turbopuffer is a fast search engine that hosts 1T+ documents, handles 10M+ writes/s, and serves 25k+ queries/s. We are ready for far more. We hope you'll trust us with your queries.

Get started