Summary

Fix OpenAI-compatible prompt_tokens_details.cached_tokens accounting for native disaggregated serving.

When a request goes through the context worker first, the generation worker sees prompt KV that was supplied by the context worker. That internal KV-transfer state was being exposed directly as OpenAI usage.prompt_tokens_details.cached_tokens, so disaggregated responses could report the entire prompt as cached even for a fresh, cache-busted request.

This PR rewrites the final client-facing usage at the disaggregated master:

• keep completion_tokens from the generation response
• take prompt_tokens and prompt_tokens_details.cached_tokens from the context response
• apply the same rewrite to streaming SSE usage events
• leave generation-only / no-context paths unchanged

Reproduction

Start a native TRT-LLM disaggregated endpoint, then send a fresh request with a unique cache-busting system prompt:

ENDPOINT=http://<host>:<port>

curl -sS --max-time 60 "$ENDPOINT/v1/chat/completions" \
  -H 'Content-Type: application/json' \
  -d "{\"model\":\"openai/gpt-oss-120b\",\"messages\":[{\"role\":\"system\",\"content\":\"You are concise. [rid:manual-$(date +%s)]\"},{\"role\":\"user\",\"content\":\"Say pong.\"}],\"temperature\":0,\"max_tokens\":32,\"stream\":false}"

Observed on a live disaggregated master:

{
  "usage": {
    "prompt_tokens": 74,
    "completion_tokens": 16,
    "total_tokens": 90,
    "prompt_tokens_details": {
      "cached_tokens": 74
    }
  }
}

The request is fresh and cache-busted, so reporting cached_tokens == prompt_tokens is not a user-facing prefix-cache hit. It is the generation worker reporting prompt KV it received from the context worker.

Fix Details

The context worker owns the external prefill/cache accounting. The generation worker owns completion accounting, but its cached-token count includes KV supplied by the context worker in the disaggregated protocol.

For requests that use a context phase, the disaggregated master now normalizes the final usage before returning it:

client prompt_tokens        <- context response usage.prompt_tokens
client cached_tokens        <- context response usage.prompt_tokens_details.cached_tokens
client completion_tokens    <- generation response usage.completion_tokens
client total_tokens         <- prompt_tokens + completion_tokens

Streaming responses are handled by rewriting only SSE events that contain a JSON usage object; content events and [DONE] pass through unchanged.

Tests

Added unit coverage for both context_first and generation_first:

• non-streaming disagg responses use context-side cached-token accounting
• streaming SSE final usage chunks are rewritten the same way

Local checks run:

python3 -m py_compile \
  tensorrt_llm/serve/openai_disagg_service.py \
  tests/unittest/disaggregated/test_openai_disagg_service.py

git diff --check

I could not run the pytest target in this sparse local checkout because the local Python environment does not have pytest installed and importing the sparse TRT-LLM package requires non-checked-out vendored runtime files.

Summary by CodeRabbit

Release Notes

• Bug Fixes
  • Fixed token usage accounting in disaggregated serving for both streaming and non-streaming modes
  • Improved accuracy of cached token tracking in token usage metrics
  • Enhanced context and generation phase token counting in disaggregated configurations