Skip to content

perf(ai): make variable-length text ingestion more efficient #394

Description

@jimezesinachi

Summary

Code-indexing workloads contain text inputs with extreme length variation: short functions, medium methods, long traits, and generated service blocks. A DevMind run against Ahnlich indexed 2,217 chunks across 190 files and took more than 40 minutes on CPU. DevMind's request shape contributes, but Ahnlich's text path should also avoid doing model work for padding that carries no information.

Related client-side findings: uche09/dev_mind#1

Primary issue: padding scope is larger than the model inference batch

The model manager preprocesses every input in an incoming request before model inference:

let processed_inputs = self.preprocess_store_input(process_action, inputs)?;
let mut store_key = self
    .model
    .model_ndarray(processed_inputs, &action_type, execution_provider, model_params)
    .await?;

For text, preprocessing tokenizes the input collection together:

let outputs = provider.preprocess_texts(inputs, truncate)?;
let token_size = outputs
    .first()
    .ok_or(/* ... */)?
    .len();

Only later does the ONNX model divide the already-prepared encodings into configured model batches:

for batch_encoding in encodings
    .into_iter()
    .chunks(self.model_batch_size)
    .into_iter()
{
    let embeddings =
        self.batch_inference_text(batch_encoding.collect(), &session)?;
}

batch_inference_text() constructs dense tensors using one encoding length for the whole inference batch:

let batch_size = encodings.len();
let encoding_length = encodings[0].len();
let max_size = encoding_length * batch_size;

for encoding in &encodings {
    ids_array.extend(encoding.get_ids().iter().map(|x| *x as i64));
}

Array::from_shape_vec((batch_size, encoding_length), ids_array)

For variable-length code, a long chunk can cause many shorter chunks to be padded and run at its sequence length. Transformer inference becomes materially more expensive as sequence length grows, so this wastes CPU and allocation pressure.

Proposed implementation

  1. Tokenize text inputs without padding the entire request together.
  2. Retain each encoding's original input position and unpadded token length.
  3. Bucket/sort similarly sized encodings into batches up to model_batch_size.
  4. Pad only within an individual model-inference batch.
  5. Run inference and restore ModelResponses to the original request order.
  6. Validate the token limit for every individual input before inference, rather than relying on the first processed encoding's length.

This should be internal to Ahnlich: callers continue to submit Set.inputs in their preferred order and receive responses in that same order.

Additional Ahnlich work exposed by the workload

Apply adaptive parallelism to AI preparation

AI ingestion uses Rayon both when assembling request data and when cloning raw strings for preprocessing:

params.inputs.into_par_iter().flat_map(/* ... */).collect()
inputs.par_iter().filter_map(/* clone RawString */).collect()

ONNX Runtime also uses CPU threads. For sustained indexing, those pools can oversubscribe the machine. Use the existing parallelism policy or equivalent request-load-aware thresholds for preparation work.

Add ingestion phase timings

Expose timings/metrics for:

  • validation and request preparation;
  • tokenization;
  • ONNX inference;
  • response/vector construction;
  • DB pipeline persistence.

Without this, a client sees only end-to-end latency and cannot establish whether model execution or storage is the limiting factor.

Make execution resources explicit

When callers omit an execution provider, text inference uses CPU. Expose documented configuration for appropriate execution providers and ONNX Runtime CPU thread settings so operators can tune ingestion to available hardware.

Future: bounded dynamic batching across requests

The per-model worker processes one request at a time. A bounded micro-batcher could coalesce concurrently arriving small requests, split responses back to callers, and improve general small-write throughput. This is a follow-up: it does not help a strictly sequential client as much as fixing padding scope and client request batching.

Acceptance criteria

  • Mixed-length text inputs preserve result count and caller input order.
  • Token-limit validation identifies the offending input deterministically.
  • Benchmarks cover representative short, medium, and long code-like inputs, not only uniform short texts.
  • The revised path demonstrates lower runtime and allocation for mixed-length batches than the current implementation.
  • Existing text-model behavior and public API remain compatible.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions