Repository navigation
feat: model response detokenization in replay - #634
Conversation
Engine replay recorded every token and every finished request at the end of the engine pass that produced it, and a closed-loop client sent its next request at that same instant. A served request also pays the server's response path before the client receives its tokens: the serving engine's server processes detokenize each pass's outputs, format streamed chunks and write them to the connections. - engine.detokenization (workers, base_ms, per_output_us) mirrors engine.tokenization. Every token-emitting worker (aggregated, or decode) has `workers` detokenizer servers (vLLM --api-server-count), shared by its attention-DP ranks, and each request belongs to one of them by a stable hash of its id. When a pass ends, each server takes the pass's outputs that belong to it as one job of base_ms + per_output_us x outputs, after its earlier jobs, and the client receives them when the job ends. Token times, latency metrics, throughput and request completion use the client's receive time, and a request that waits on a finished response is released when the client receives its end. - A rejection carries no output: it skips detokenization and reaches the client when the engine rejects the request, in both topologies. - The SGLang host-loop TTFT breakdown adds mean_detokenization_ms, the first token's emission to its receipt by the client (zero without detokenization), so the stages still sum to TTFT. - Off by default and keeps outputs request for request. Offline token replay on --stack engine only; worker scaling, AFD, encoder pools and agentic profiles are rejected. Rust: Replayer::with_detokenization. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yongming Ding <yongmingd@nvidia.com>
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Detokenization model and replay API crates/core/src/replay/components/detokenization.rs, crates/core/src/replay/components/mod.rs, crates/core/src/replay/replayer.rs, crates/core/tests/replay_detokenization.rs, docs/replay/api/rust.md |
The replay API configures detokenization latency. The stage assigns requests to servers, queues pass output jobs, and releases terminals by receipt time. Tests cover timing, queueing, validation, and zero-cost settings. |
Python configuration and execution wiring python/aisimulate/src/aisimulate/config/engine.py, python/aisimulate/src/aisimulate/config/cli.py, python/aisimulate/src/aisimulate/compiler.py, python/aisimulate/src/aisimulate/runner.py, python/aisimulate/src/aisimulate/sweeper/replay.py, python/aisimulate/src/aisimulate/sweeper/score.py, crates/core/src/python.rs, tests/test_runner.py, docs/replay/engine/README.md |
Python config validates detokenization settings, compiles active settings into deployment fields, and passes them to compatible offline replay runners. The execution payload accepts settings for aggregated and disaggregated topologies. |
Runtime receipt scheduling and delivery crates/core/src/replay/agg.rs, crates/core/src/replay/disagg.rs, crates/core/src/replay/disagg_tests.rs, crates/core/tests/replay_detokenization.rs |
Both runtimes schedule pass outputs and terminal releases at client receipt times. Pending deliveries count as in-flight requests. Rejections bypass detokenization. Tests cover delivery timing, admission progression, and simulation cutoffs. |
TTFT reporting and validation crates/core/src/replay/report.rs, crates/core/src/replay/report/bounded.rs, crates/core/tests/engine.rs, python/aisimulate/src/aisimulate/sweeper/score.py, docs/getting-started/understand-results.md, docs/reference/cli.md, docs/replay/engine/sglang-vl-host-loop.md |
TTFT accounting and serialization add a detokenization span. Tests and documentation describe the added metric and client-receipt timestamps. |
Priority:
Estimated code review effort: 4 (Complex) | ~45 minutes
Possibly related PRs
- ai-dynamo/aisimulate#624: Changes vLLM admission at pass boundaries, which interacts with detokenization-delayed arrivals and terminal release.
- ai-dynamo/aisimulate#612: Changes mixed prefill/decode pass latency, which affects the pass completion times used to schedule detokenization.
Merge Risk: 
86795
-
With detokenization enabled and a simulation time cap that falls between a pass end and its client receipt, an incomplete request can report a first-token time after the cap. Completion counts and closed-loop release are unaffected. This is a bounded reporting edge case, acceptable as a follow-up.
Pre-merge checks: 
-
Check name Status Explanation Resolution Cross-Layer Contract The prior cross-layer failures remain unresolved. docs/replay/features.md:37is unchanged and still states that native SGLang VL replay has no output-side costs, while this PR adds `engine.detokeniz…Update docs/replay/features.mdfor optional response detokenization in SGLang host-loop replay. Enforce the documented encoder-pool exclusion at the RustReplayer/JSON boundary, or remove the exclusion consistently from Python and docum…Compatibility Boundaries The new Python and Rust defaults are not fully synchronized for inactive detokenization. Python treats engine.detokenization: {}or all-zero costs as inactive (DetokenizationConfig.activeand comp…Normalize or test detokenization activity before the Rust agentic-profile rejection. Use base_ms > 0.0 || per_output_ms > 0.0for the compatibility check, while still passing the block towith_detokenizationso invalid workers and non-f…Modeling And Data Evidence The available evidence does not establish the required modeling/data evidence. The PR adds deterministic formula tests that compare direct and detokenized replay, such as `crates/core/tests/replay_det… Provide the complete PR description or commit reproducible evidence. Include the source and version for the detokenization model, raw machine-readable before/after results, the comparison procedure, and anomaly checks. Unit assertions alone… Review Evidence The visible PR description names commands and results, identifies a Linux VM run, reports parity with main, and states several rejection boundaries. However, the description is truncated in the eviden… Provide the complete PR description, including the remainder of the Evidence section and all CI/evidence labels. Then verify that each result is explicitly classified as local, fake-runner, parity, or production evidence and that no inactiv… -
Check name Status Explanation Title check The title "feat: model response detokenization in replay" states the main behavioral change: replay now models response detokenization. The changeset adds an opt-in engine.detokenization timing model,… Description check The description follows the template: it covers why and what changed, a review map with risk level, start-with files, and the contract change, plus evidence (commands, test counts, commit, negative ca… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Full details: Title check
-
Explanation
The title "feat: model response detokenization in replay" states the main behavioral change: replay now models response detokenization. The changeset adds an opt-in engine.detokenization timing model, DetokenizationLatency/DetokenizationStage, and the mean_detokenization_ms TTFT stage, all matching the title. The title follows the conventional prefix form and is specific. It does not state the 'precise' behavioral effect (client-receipt timing), but it is related to the main change and is not vague, so it passes.
Full details: Description check
-
Explanation
The description follows the template: it covers why and what changed, a review map with risk level, start-with files, and the contract change, plus evidence (commands, test counts, commit, negative cases, and calibration data), and a modeling provenance section with units and limitations. Tracking is filled as N/A and related PRs. Two template gaps remain. 'Fast CI' and 'Full CI' are blank. 'CodeRabbit reviewed commit' is missing. 'Codex reviewed commit' is also missing, and the template marks it as required for medium or high risk. The author labels risk as medium, so that field is required. 'Expected-value derivation' is not explicitly addressed for the exact numerical assertions in the new tests.
Resolution
Fill in the blank template fields before merge. Add the Fast CI and Full CI results. Add the CodeRabbit reviewed commit. Add the Codex reviewed commit, which is required for medium-risk changes. Add the expected-value derivation for exact numerical assertions, or state N/A if none apply.
-
Full details: Cross-Layer Contract
-
Explanation
The prior cross-layer failures remain unresolved.
docs/replay/features.md:37is unchanged and still states that native SGLang VL replay has no output-side costs, while this PR addsengine.detokenizationand SGLang detokenization tests (crates/core/tests/engine.rs:2239-2288). The documented encoder-pool exclusion is not enforced by the Rust API: Python rejects active detokenization with an encoder (python/aisimulate/src/aisimulate/config/engine.py:917-922), butReplayer::run_inneronly checks scaling (crates/core/src/replay/replayer.rs:531-537) and passes both encoder and detokenization to the runtimes (:598-600,:676-678).RunnerCapabilities.require_compatiblechecks only AFD, capability, and offline mode for detokenization (python/aisimulate/src/aisimulate/sweeper/replay.py:572-578); it has no detokenization-specific scaling or agentic-profile checks, and the added tests cover neither combination.BackendDeploymentSpec.detokenizationis marked for canonical omission (python/aisimulate/src/aisimulate/sweeper/replay.py:212-214), but the only canonical JSON test still covers tokenization (tests/test_runner.py:2966-2973).Resolution
Update
docs/replay/features.mdfor optional response detokenization in SGLang host-loop replay. Enforce the documented encoder-pool exclusion at the RustReplayer/JSON boundary, or remove the exclusion consistently from Python and documentation, and add direct tests. AddRunnerCapabilitiespreflight rejection and tests for detokenization combined with scaling andagentic_profile. Add canonical JSON tests proving configured detokenization is serialized and unset detokenization is omitted.
Full details: Compatibility Boundaries
-
Explanation
The new Python and Rust defaults are not fully synchronized for inactive detokenization. Python treats
engine.detokenization: {}or all-zero costs as inactive (DetokenizationConfig.activeand compiler omission atpython/aisimulate/src/aisimulate/config/engine.py:889-890,python/aisimulate/src/aisimulate/compiler.py:215-223). The Rust binding accepts the same defaulted payload, but rejects it with anagentic_profilebecause it checksserving_path.detokenization.is_some()beforeReplayer::with_detokenizationremoves zero-cost latency (crates/core/src/python.rs:2950-2955,crates/core/src/replay/replayer.rs:446-456). This contradicts the documented default that zero costs are equivalent to omission (docs/replay/engine/README.md:237-241) and changes the new public execution boundary. Version pins remain equal at 0.13.0, the release scripts still enforce oneaisimulatewheel and oneaisimulate-corecrate, and no Dynamo dependency was added.Resolution
Normalize or test detokenization activity before the Rust agentic-profile rejection. Use
base_ms > 0.0 || per_output_ms > 0.0for the compatibility check, while still passing the block towith_detokenizationso invalid workers and non-finite costs remain validated. Add a binding test for a defaulted or all-zero detokenization payload withagentic_profileand verify it behaves like an omitted block.
Full details: Modeling And Data Evidence
-
Explanation
The available evidence does not establish the required modeling/data evidence. The PR adds deterministic formula tests that compare direct and detokenized replay, such as
crates/core/tests/replay_detokenization.rs:116-239andcrates/core/tests/engine.rs:2219-2292, but the changed files contain no explained golden diff, independent oracle, held-out comparison, or machine-readable results with anomaly checks. The documentation explicitly says the example timing values are illustrative and must be measured on each deployment (docs/replay/engine/README.md:483-486). The authored PR description is truncated after its evidence section, so omitted provenance or reproducible before/after data cannot be assessed.Resolution
Provide the complete PR description or commit reproducible evidence. Include the source and version for the detokenization model, raw machine-readable before/after results, the comparison procedure, and anomaly checks. Unit assertions alone do not establish empirical model or performance evidence.
Full details: Review Evidence
-
Explanation
The visible PR description names commands and results, identifies a Linux VM run, reports parity with main, and states several rejection boundaries. However, the description is truncated in the evidence section after “New tests,” so the required distinction between local, fake-runner, parity, and production evidence cannot be verified. The omitted text may also contain or omit hosted-CI claims; the visible blank “Fast CI” and “Full CI” headings do not resolve that question.
Resolution
Provide the complete PR description, including the remainder of the Evidence section and all CI/evidence labels. Then verify that each result is explicitly classified as local, fake-runner, parity, or production evidence and that no inactive nested workflow is presented as hosted CI evidence.
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
tedzhouhk
left a comment
tedzhouhk
left a comment
There was a problem hiding this comment.
Reviewed 4f6264f across the aggregated and disaggregated client response paths. No blocking issues found. Local validation: 13 targeted Rust tests and 17 Python tests passed, plus independent probes covering 78 combinations of slow output queues, bounded/detailed report parity, closed-loop concurrency, and completion dependencies. GPU accuracy experiments were not rerun.
Rename client_path.rs to detokenization.rs and ClientPath to DetokenizationStage, mirroring tokenization.rs, and rename the agg and disagg fields, helpers and tests that still said client path. No behavior change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yongming Ding <yongmingd@nvidia.com>
|
/ok to test 86795be |
Signed-off-by: Yongming Ding <yongmingd@nvidia.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Why and what changed
Engine replay recorded every token and every finished request at the end of the engine pass that produced it, and a closed-loop client sent its next request at that same instant. A served request also pays the server's response path before the client receives its tokens: the serving engine's server processes detokenize each pass's outputs, format streamed chunks and write them to the connections. We found this gap while running the paper's vLLM replay experiments: replay under-predicted TTFT most at high concurrency, where many outputs leave the same pass.
engine.detokenization(workers,base_ms,per_output_us) mirrorsengine.tokenization. Every token-emitting worker (aggregated, or decode in P/D) hasworkersdetokenizer servers, shared by its attention-DP ranks; a request belongs to one of them by a stable hash of its id. When a pass ends, each server takes that pass's outputs that belong to it as one FIFO job ofbase_ms + per_output_us × outputs, and the client receives them when the job ends.AsyncLLM._run_output_handler) that processes one engine step's outputs at a time for the requests it owns, whichever DP rank produced them.workersmaps to--api-server-count.mean_detokenization_ms(first token emitted → received by the client; zero without detokenization), so the stages still sum to TTFT.--stack engineonly; worker scaling, AFD, encoder pools and agentic profiles are rejected. Rust:Replayer::with_detokenization.Review map
crates/core/src/replay/components/detokenization.rs(DetokenizationLatency,DetokenizationStage::receive_pass), thencrates/core/src/replay/agg.rs(receive_pass,process_output_signal,release_received_terminals) andcrates/core/src/replay/disagg.rs(receive_decode_pass,notify_received_terminal,awaiting_receiptin quiescence), thencrates/core/tests/replay_detokenization.rs.engine.detokenizationYAML block,Replayer::with_detokenization,mean_detokenization_msin the TTFT stage breakdown (TTFT_STAGE_COUNT6 → 7, always present, zero without detokenization).Evidence
cargo fmt --all -- --check;cargo test --workspace(core 2,369 passed) andcargo test --workspace --features embed-python,replay-benchwithAIC_REQUIRE_EMBEDDED_ROUND_TRIP=1(core 2,457 passed) pass with no failures;cargo test --manifest-path crates/tests/public-api/Cargo.tomlpasses (21).ruff check,ruff format --checkandscripts/ci/check_packaged_legal_files.pypass. Pythontests/andpython/aisimulate/tests -m unitfail exactly the same tests as main on that host (204 and 72; they fail on main too) and run 17 new cases.crates/core/tests/replay_detokenization.rs(aggregated and disaggregated: a lone request's outputs after detokenization, outputs of one pass reach the client together, the closed-loop client sends at receive time, rejections reach the client at engine time and do not delay their pass's outputs, attention-DP ranks share a worker's servers,base_msoccupies the detokenizer, zero costs keep pass-end delivery, the time cap, agentic requests, invalid configs),sglang_ttft_stages_attribute_detokenization_to_its_own_stage,DetokenizationStageunit tests (two servers split a pass and queue independently, a request keeps its server, a pass without outputs takes no server time, atomic rejection of a non-finite job end), Python compiler/config/runner cases and a native end-to-end replay.base_ms2.85,per_output_us61.6. Mean absolute error per shape is 5.6 ms for the first token and 2.1 ms for the response end; a four-term model (separate first-output and final-output costs) fits 1.1 and 0.9 ms better, and we kept the two-term model for simplicity. No evaluation point was used.engine.detokenization:workers: 1counts outputs without allocating); nothing without it. Nativewall_time_ms, release builds of f45d14a (main) and this PR on an idle VM, paired alternating runs. Off: Exp1 1k c1 / 1k c256 / 8k c64 median ratio to main 1.035 / 0.994 / 0.992 (5 runs). Synthetic 512/128 closed loops of 1k / 4k / 16k and a 400 req/s Poisson overload (3 runs, every request completed, same simulated outcome across variants), host cost per output token main → detokenization (workers: 1) →workers: 4: 0.49 → 0.54 → 0.57, 0.66 → 0.65 → 0.69, 0.72 → 0.67 → 0.70, 0.72 → 0.72 → 0.71 µs. Theworkers: 4path allocates two small vectors per pass; its c1024 ratio is 1.17 [1.00, 1.34] and the larger runs are within noise.engine.detokenization. Fixed: TTFT stage attribution, agg/disagg rejection symmetry, all-zero costs equal omission, unknown rank is an error. Documented: per-pass delivery vs vLLM's 128-output chunks,--api-server-countdefaults.Modeling or data provenance
VLLM_V1_OUTPUT_PROC_CHUNK_SIZE(128 by default), so with more than 128 outputs per server in one pass, replay delivers the earlier chunks somewhat late.Tracking
🤖 Generated with Claude Code