Repository navigation
fix(server): report the expanded media prompt length on the xla-iree loop #2420
Description
Problem
PR #2415 (closed #2337) holds a media body's streaming message_start.usage.input_tokens until the worker's first GenerateEvent::Prefill (src/server/streaming_anthropic.rs:311 AnthropicMessageStart), and docs/llama-server-compat.md:129 promises that media counts are the expanded soft-token length. The xla-iree serve loop never sends Prefill and reports the pre-expansion length everywhere, so for LLaVA and Phi4MM media bodies message_start, usage and /v1/messages/count_tokens all carry the placeholder count.
Evidence (main 757c0ba)
src/server/batch/xla_worker.rs:434-520dispatchsends onlyToken,DoneandError.src/server/batch/xla_worker_admission.rs:419and:481setprompt_token_count = state.prompt_tokens.len()(logical) next toeffective_prefill_len = prepared.sequence_len();xla_worker.rs:509,545put the logical count in the result. The XLA worker ignorescount_prompt_only, so the count route (src/server/model_provider_prompt_count.rs:65) reports the same figure.xla_worker_tests.rs:450pins this: the image request expands to 5 tokens and assertsprompt_tokens == 3at:499.
Fix
Make effective_prefill_len the public count on this loop, reversing the #895 split, so the XLA path matches the batch scheduler and the docs. Add prefill_reported: bool to ServeState (xla_worker.rs:237); in dispatch, before a request's first Token or its Finished, send PrefillStats { prompt_tokens: effective_prefill_len, cached_tokens: 0, processed: effective_prefill_len, first_token: true, prompt_ms }, as src/server/batch/sequence.rs:1200-1219 does. Pass the effective length to detok.finish/finish_truncated and to the finish_zero_budget calls at xla_worker_admission.rs:414,477. Text requests are unchanged (both lengths are equal).
Out of scope
The disaggregated coordinator: coordinator.rs:906 drops Prefill, but the router front serves only /v1/chat/completions and /v1/completions (src/server/router_front.rs:2062-2065) and refuses media (:885, :945), so no media or Anthropic body reaches it.
Acceptance Criteria
- Driven through the real
XlaServeWorker::admit, preprocessing anddispatchpath, an image and an audio request each receivePrefillwith the expanded length before their firstToken, and theirDonereports that length; the updated:450test and a new ordering test fail on current main. - A text request's
Prefillequals its prompt length.
Verification
cargo check --features cuda,xla-iree --all-targets
cargo test --release --features cuda,xla-iree --lib server::batch::xla_worker -- --test-threads=1
Activity
fix(server): report the expanded media prompt length on the xla-iree …
fix(server): report the expanded media prompt length on the xla-iree …