Skip to content

fix(server): report the expanded media prompt length on the xla-iree loop #2420

Closed

Description

@inureyes

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-520 dispatch sends only Token, Done and Error.
  • src/server/batch/xla_worker_admission.rs:419 and :481 set prompt_token_count = state.prompt_tokens.len() (logical) next to effective_prefill_len = prepared.sequence_len(); xla_worker.rs:509,545 put the logical count in the result. The XLA worker ignores count_prompt_only, so the count route (src/server/model_provider_prompt_count.rs:65) reports the same figure.
  • xla_worker_tests.rs:450 pins this: the image request expands to 5 tokens and asserts prompt_tokens == 3 at :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 and dispatch path, an image and an audio request each receive Prefill with the expanded length before their first Token, and their Done reports that length; the updated :450 test and a new ordering test fail on current main.
  • A text request's Prefill equals 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

  1. added
    type:bugBug fixes, error corrections, or issue resolutions
    area:inferenceGeneration, sampling, decoding (incl. speculative, DRY)
    and removed on Oct 16, 2026
  2. added a commit that references this issue on Oct 17, 2026
    815a6da
  3. added a commit that references this issue on Oct 18, 2026
    85c1fe7
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:inferenceGeneration, sampling, decoding (incl. speculative, DRY)priority:lowLow prioritystatus:doneCompletedtype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions