[None][perf] disagg: skip per-chunk JSON parse on streaming usage rewrite by lancelly · Pull Request #15042 · NVIDIA/TensorRT-LLM

lancelly · 2026-06-06T13:38:48Z

Description

In disaggregated serving the orchestrator relays the generation server's streaming
SSE response back to the client, and rewrites the usage object so the reported
prompt_tokens / cached_tokens reflect the context server (where prefill
actually happened) rather than the generation server.

_rewrite_usage_sse_event_from_ctx was invoked for every relayed SSE chunk and
ran splitlines() + json.loads() on each one — but a top-level usage object
only appears on the final chunk (and only when the client requested usage).
So for essentially every streamed output token the orchestrator did a full JSON
parse just to discover "no usage here, nothing to rewrite". On the orchestrator's
single asyncio event loop this is pure per-token overhead that serializes all
concurrent streams: ~5% of one core at high concurrency (~15k relayed tok/s).

This PR adds a cheap byte pre-filter before the parse:

if b'"usage"' not in event:
    return event

A top-level usage field cannot exist unless the chunk bytes contain "usage",
so chunks without it skip the splitlines() + json.loads() entirely and are
returned verbatim — byte-for-byte identical to the previous no-usage behavior.
The rewrite is never dropped (when a usage field exists, "usage" is present in
the bytes, so the chunk still takes the full path); a rare token whose content is
literally "usage" simply falls through and re-parses harmlessly (the nested
match is not a top-level key, so it returns unchanged).

Measured (json.dumps-shaped chat.completion.chunk, in-container aarch64):
_rewrite_usage_sse_event_from_ctx drops from 3.46 µs → 0.25 µs per chunk
(~5% → ~0.4% of one core at 15.4k tok/s). Behavior is unchanged; this only avoids
work on chunks that have nothing to rewrite.

…rite When rewriting a disaggregated streaming response's usage from the context server, _rewrite_usage_sse_event_from_ctx ran json.loads on EVERY relayed SSE chunk, even though a top-level usage object only appears on the final chunk (and only when the client requested usage). On the orchestrator's single event loop this full JSON parse of every relayed output token is ~5% of one core at high concurrency (~15k tok/s). Add a byte pre-filter: a top-level usage field cannot exist unless the chunk bytes contain "usage", so skip the splitlines + json.loads for chunks that don't. The event is returned verbatim on the fast path (identical to the prior no-usage behavior), so the rewrite is never dropped; a rare token whose content is literally "usage" just re-parses harmlessly. Micro-bench: 3.46us -> 0.25us per chunk (~5%% -> ~0.4%% of one core). Existing streaming usage-rewrite unit test (test_send_disagg_request) covers both paths and passes. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>

lancelly · 2026-06-08T06:48:20Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-08T06:54:06Z

PR_Github #52691 [ run ] triggered by Bot. Commit: 2ce78b4 Link to invocation

tensorrt-cicd · 2026-06-08T08:27:59Z

PR_Github #52691 [ run ] completed with state FAILURE. Commit: 2ce78b4
/LLM/main/L0_MergeRequest_PR pipeline #41958 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lancelly · 2026-06-08T08:30:30Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-08T08:36:12Z

PR_Github #52721 [ run ] triggered by Bot. Commit: 2ce78b4 Link to invocation

tensorrt-cicd · 2026-06-08T11:23:54Z

PR_Github #52721 [ run ] completed with state SUCCESS. Commit: 2ce78b4
/LLM/main/L0_MergeRequest_PR pipeline #41987 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lancelly · 2026-06-08T12:14:11Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-08T12:20:03Z

PR_Github #52768 [ run ] triggered by Bot. Commit: 2ce78b4 Link to invocation

tensorrt-cicd · 2026-06-08T14:00:26Z

PR_Github #52768 [ run ] completed with state SUCCESS. Commit: 2ce78b4
/LLM/main/L0_MergeRequest_PR pipeline #42028 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

peihu-nv · 2026-06-08T16:59:19Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-08T17:05:42Z

PR_Github #52822 [ run ] triggered by Bot. Commit: 2ce78b4 Link to invocation

tensorrt-cicd · 2026-06-08T18:44:36Z

PR_Github #52822 [ run ] completed with state SUCCESS. Commit: 2ce78b4
/LLM/main/L0_MergeRequest_PR pipeline #42075 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

peihu-nv · 2026-06-08T19:09:54Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-08T19:16:05Z

PR_Github #52848 [ run ] triggered by Bot. Commit: 2ce78b4 Link to invocation

tensorrt-cicd · 2026-06-08T20:54:24Z

PR_Github #52848 [ run ] completed with state SUCCESS. Commit: 2ce78b4
/LLM/main/L0_MergeRequest_PR pipeline #42101 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

peihu-nv

CI failure looks to be flaky test. Not due to the PR

…rite (NVIDIA#15042) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>

…VIDIA#15042 NVIDIA#15042 ("skip per-chunk JSON parse on streaming usage rewrite") was a 5-line optimization on feat/deepseek_v4, layered on the method-based usage-rewrite machinery introduced by NVIDIA#13620. When this PR carved NVIDIA#15042 onto main, that machinery did not exist there: main's NVIDIA#14177 ("Fix disaggregated usage propagation") had refactored NVIDIA#13620 into postprocess_handlers (rewrite_usage_info_from_ctx, operating on parsed chunk objects). The carve folded the obsolete pre-refactor machinery back in to make the patch apply. The result is dead code: _rewrite_disagg_usage (the chain root) has no caller, so the whole _ctx_usage_for_client / _rewrite_usage_* / _rewrite_streaming_usage_from_ctx subtree is unreachable, and NVIDIA#15042's json.loads-skip optimization sits inside it (a no-op). main already rewrites disagg usage via postprocess_handlers and relays streaming bytes untouched, so the optimized per-chunk-parse path no longer exists. Remove the dead methods and the now-orphaned imports (json, AsyncIterator, PromptTokensDetails). No behavior change: the live usage path is unaffected. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>

…rite (NVIDIA#15042) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>

…VIDIA#15042 NVIDIA#15042 ("skip per-chunk JSON parse on streaming usage rewrite") was a 5-line optimization on feat/deepseek_v4, layered on the method-based usage-rewrite machinery introduced by NVIDIA#13620. When this PR carved NVIDIA#15042 onto main, that machinery did not exist there: main's NVIDIA#14177 ("Fix disaggregated usage propagation") had refactored NVIDIA#13620 into postprocess_handlers (rewrite_usage_info_from_ctx, operating on parsed chunk objects). The carve folded the obsolete pre-refactor machinery back in to make the patch apply. The result is dead code: _rewrite_disagg_usage (the chain root) has no caller, so the whole _ctx_usage_for_client / _rewrite_usage_* / _rewrite_streaming_usage_from_ctx subtree is unreachable, and NVIDIA#15042's json.loads-skip optimization sits inside it (a no-op). main already rewrites disagg usage via postprocess_handlers and relays streaming bytes untouched, so the optimized per-chunk-parse path no longer exists. Remove the dead methods and the now-orphaned imports (json, AsyncIterator, PromptTokensDetails). No behavior change: the live usage path is unaffected. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>

…rite (NVIDIA#15042) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com> Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com>

…VIDIA#15042 NVIDIA#15042 ("skip per-chunk JSON parse on streaming usage rewrite") was a 5-line optimization on feat/deepseek_v4, layered on the method-based usage-rewrite machinery introduced by NVIDIA#13620. When this PR carved NVIDIA#15042 onto main, that machinery did not exist there: main's NVIDIA#14177 ("Fix disaggregated usage propagation") had refactored NVIDIA#13620 into postprocess_handlers (rewrite_usage_info_from_ctx, operating on parsed chunk objects). The carve folded the obsolete pre-refactor machinery back in to make the patch apply. The result is dead code: _rewrite_disagg_usage (the chain root) has no caller, so the whole _ctx_usage_for_client / _rewrite_usage_* / _rewrite_streaming_usage_from_ctx subtree is unreachable, and NVIDIA#15042's json.loads-skip optimization sits inside it (a no-op). main already rewrites disagg usage via postprocess_handlers and relays streaming bytes untouched, so the optimized per-chunk-parse path no longer exists. Remove the dead methods and the now-orphaned imports (json, AsyncIterator, PromptTokensDetails). No behavior change: the live usage path is unaffected. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>

…rite (NVIDIA#15042) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com> Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com>

…VIDIA#15042 NVIDIA#15042 ("skip per-chunk JSON parse on streaming usage rewrite") was a 5-line optimization on feat/deepseek_v4, layered on the method-based usage-rewrite machinery introduced by NVIDIA#13620. When this PR carved NVIDIA#15042 onto main, that machinery did not exist there: main's NVIDIA#14177 ("Fix disaggregated usage propagation") had refactored NVIDIA#13620 into postprocess_handlers (rewrite_usage_info_from_ctx, operating on parsed chunk objects). The carve folded the obsolete pre-refactor machinery back in to make the patch apply. The result is dead code: _rewrite_disagg_usage (the chain root) has no caller, so the whole _ctx_usage_for_client / _rewrite_usage_* / _rewrite_streaming_usage_from_ctx subtree is unreachable, and NVIDIA#15042's json.loads-skip optimization sits inside it (a no-op). main already rewrites disagg usage via postprocess_handlers and relays streaming bytes untouched, so the optimized per-chunk-parse path no longer exists. Remove the dead methods and the now-orphaned imports (json, AsyncIterator, PromptTokensDetails). No behavior change: the live usage path is unaffected. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>

github-actions Bot assigned lancelly Jun 6, 2026

lancelly force-pushed the perf/disagg-skip-sse-usage-parse branch from 7ab5ff8 to 641f858 Compare June 6, 2026 13:41

lancelly force-pushed the perf/disagg-skip-sse-usage-parse branch from 641f858 to 2ce78b4 Compare June 6, 2026 13:45

lancelly marked this pull request as ready for review June 8, 2026 06:48

lancelly requested a review from a team as a code owner June 8, 2026 06:48

lancelly requested review from zhenhuaw-me and removed request for a team June 8, 2026 06:48

peihu-nv approved these changes Jun 8, 2026

View reviewed changes

peihu-nv merged commit eba87a6 into NVIDIA:feat/deepseek_v4 Jun 8, 2026
6 of 7 checks passed

Uh oh!

Conversation

lancelly commented Jun 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Uh oh!

lancelly commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

lancelly commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

lancelly commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

peihu-nv commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

peihu-nv commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

peihu-nv left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

lancelly commented Jun 6, 2026 •

edited

Loading