fix(frontend): offload prompt tokenization off the async event loop

[nv-yna]

What this PR does

Moves prompt tokenization off the frontend's async (tokio) event loop and onto the bounded blocking
thread pool.

Why / the improvement

The generate path tokenizes the prompt synchronously on the async event loop
(gather_tokens → encode_with_timing → tokenizer.encode), unlike the embedding path
(preprocess_embedding_request), which already offloads encoding via spawn_blocking. For long prompts
under concurrency, the CPU-heavy BPE encode blocks the runtime for multiple seconds, stalling every other
task that shares it.

Offloading the encode keeps the event loop responsive while a prompt is tokenized, so request latency and
throughput no longer degrade when large prompts are encoded under load. In disaggregated serving this also
stops the stall from starving request-plane I/O (which could otherwise cause healthy workers to be wrongly
marked down).

How

• encode_with_timing becomes async; tokenizer.encode runs inside tokio::task::spawn_blocking.
• The prompt is owned and the tokenizer Arc is cloned so the closure is 'static + Send.
• async is propagated up through gather_tokens and its call sites. The one synchronous FFI consumer
  (lib/bindings/c EPP tokenize path) bridges the now-async call via the crate's existing
  runtime.secondary().block_on(...).
• Mirrors the existing embedding-path offload and produces the same Encoding; null-byte stripping is
  preserved.

Scope

2 files — lib/llm/src/preprocessor.rs and lib/bindings/c/src/lib.rs (+24/−12).

Testing

cargo fmt -- --check and cargo clippy --no-deps --all-targets -- -D warnings (full workspace) both pass
on top of main.

[datadog-official]

🎯 Code Coverage (details)
• Patch Coverage: 100.00%
• Overall Coverage: 41.51% (-3.89%)

_{This comment will be updated automatically if new data arrives.
🔗 Commit SHA: 625ea9b | Docs | Give us feedback!}

[coderabbitai]

Walkthrough

The preprocessor's tokenization path is converted from synchronous to asynchronous. gather_tokens and encode_with_timing are now async, with encode_with_timing offloading tokenizer encoding via tokio::task::spawn_blocking. All internal call sites are updated to await these methods.

Changes

Async Tokenization Conversion

Layer / File(s)	Summary
Async encode_with_timing implementation lib/llm/src/preprocessor.rs	encode_with_timing is converted to async, strips null bytes using an owned string, clones the tokenizer, and performs tokenizer.encode inside tokio::task::spawn_blocking, while still recording tokenize latency via RequestTracker; the unused Cow import is removed.
gather_tokens made async and call sites updated lib/llm/src/preprocessor.rs	gather_tokens signature changes to async fn, and callers in preprocess_request_with_options, completion preprocessing, and both PromptInput::Text (Single/Batch) branches now await gather_tokens or encode_with_timing.

Estimated code review effort: 2 (Simple) | ~10 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Preprocessor
  participant gather_tokens
  participant encode_with_timing
  participant TokioBlockingPool
  participant RequestTracker

  Preprocessor->>gather_tokens: await gather_tokens(request)
  gather_tokens->>encode_with_timing: await encode_with_timing(prompt)
  encode_with_timing->>TokioBlockingPool: spawn_blocking(tokenizer.encode)
  TokioBlockingPool-->>encode_with_timing: token ids
  encode_with_timing->>RequestTracker: record tokenize latency
  encode_with_timing-->>gather_tokens: tokens
  gather_tokens-->>Preprocessor: tokens, metadata

Loading

Related issues: None referenced.

Related PRs: None referenced.

Suggested labels: performance, rust, preprocessing

Suggested reviewers: None specified.

🐰 A hop, a skip, a tokenizer's dance,
Async now gives blocking a chance,
No more waiting on the main thread's beat,
Tokens spawn off, light on their feet!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	It covers the change, rationale, scope, and testing, but it does not follow the required template or include the mandatory Related Issues section.	Rewrite the PR description to use the template headings and add the required Related Issues section with either Closes #XXXX or the no-issue checkbox path.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main change: moving prompt tokenization off the async event loop.

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

lib/llm/src/preprocessor.rs (1)

1694-1712: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consolidate duplicate encode_with_timing calls.

Both the has_backend_instance_id and fallback branches now perform the identical self.encode_with_timing(prompt, tracker).await? followed by the same tuple construction — only the warning log differs. Since both branches were touched to add .await, this is a good opportunity to de-duplicate.

♻️ Proposed refactor

-                            let (tokens_vec, skip_token_annotation) = if let Some(tokens) =
-                                token_data
-                            {
+                            if has_backend_instance_id && token_data.is_none() {
+                                tracing::warn!(
+                                    "backend_instance_id provided but no token_data; tokenizing prompt"
+                                );
+                            }
+                            let (tokens_vec, skip_token_annotation) = if let Some(tokens) =
+                                token_data
+                            {
                                 tracing::info!(
                                     token_count = tokens.len(),
                                     first_tokens = ?&tokens[..std::cmp::min(5, tokens.len())],
                                     "[SIDECAR-SKIP-TOKENIZE] Found nvext.token_data — using pre-computed tokens, SKIPPING tokenization"
                                 );
                                 (tokens.clone(), true)
-                            } else if has_backend_instance_id {
-                                tracing::warn!(
-                                    "backend_instance_id provided but no token_data; tokenizing prompt"
-                                );
-                                let encoding = self.encode_with_timing(prompt, tracker).await?;
-                                (encoding.token_ids().to_vec(), false)
                             } else {
                                 let encoding = self.encode_with_timing(prompt, tracker).await?;
                                 (encoding.token_ids().to_vec(), false)
                             };

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/llm/src/preprocessor.rs` around lines 1694 - 1712, The
`encode_with_timing` call in `preprocessor.rs` is duplicated across the
`has_backend_instance_id` and fallback paths in the `tokens_vec` /
`skip_token_annotation` branch. Refactor the `token_data` handling inside
`preprocessor.rs` so `self.encode_with_timing(prompt, tracker).await?` and the
`(token_ids, false)` tuple are computed once, while preserving the
`tracing::warn!` only in the `has_backend_instance_id` case and keeping the
`tracing::info!` path for precomputed `token_data`.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@lib/llm/src/preprocessor.rs`:
- Around line 1694-1712: The `encode_with_timing` call in `preprocessor.rs` is
duplicated across the `has_backend_instance_id` and fallback paths in the
`tokens_vec` / `skip_token_annotation` branch. Refactor the `token_data`
handling inside `preprocessor.rs` so `self.encode_with_timing(prompt,
tracker).await?` and the `(token_ids, false)` tuple are computed once, while
preserving the `tracing::warn!` only in the `has_backend_instance_id` case and
keeping the `tracing::info!` path for precomputed `token_data`.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e61b9d12-3077-4961-a838-d33c62b2b480

📥 Commits

Reviewing files that changed from the base of the PR and between 1ef2f94 and 53feb96.

📒 Files selected for processing (1)

• lib/llm/src/preprocessor.rs

[dynamo-review-agent]

lib/bindings/c/src/lib.rs — This call still matches on gather_tokens as if it returned Result, but the PR changed it to return a future, so the C bindings no longer compile. Fix: run the async tokenization on the binding runtime and await the result before matching.

🤖 AI Fix

In lib/bindings/c/src/lib.rs, inside preprocess_request, wrap the completion-branch preprocessor.gather_tokens(&request, None, None) call in handles.runtime.secondary().block_on(async { preprocessor.gather_tokens(&request, None, None).await }) and keep matching on that returned Result.