[https://nvbugs/6373447][fix] Restore rc17 throughput by trimming per-iter host overhead#16047
Closed
chenfeiz0326 wants to merge 1 commit into
Closed
[https://nvbugs/6373447][fix] Restore rc17 throughput by trimming per-iter host overhead#16047chenfeiz0326 wants to merge 1 commit into
chenfeiz0326 wants to merge 1 commit into
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughModifies PyTorchModelEngine's input preparation to remove cached_kv_tokens tracking, inline runtime_tokens_per_gen_step for non-spec-decode paths, and snapshot per-request attributes into locals for reduced repeated reads. Also gates PyExecutor's model_engine input copy wait calls behind an _is_kv_manager_v2 check. ChangesModel engine input preparation optimizations
Estimated code review effort: 3 (Moderate) | ~20 minutes Gate Model Engine Input Copy Wait Behind KV Manager V2
Estimated code review effort: 2 (Simple) | ~10 minutes Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant ModelEngine
participant KVCacheManager
PyExecutor->>PyExecutor: check _is_kv_manager_v2
alt is_kv_manager_v2 true
PyExecutor->>ModelEngine: wait_for_model_engine_input_copy()
ModelEngine-->>PyExecutor: input copy complete
else is_kv_manager_v2 false
PyExecutor->>PyExecutor: skip wait
end
PyExecutor->>KVCacheManager: _update_request_states / _process_previous_batch
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
yechank-nvidia
approved these changes
Jul 7, 2026
| self.iter_states['num_ctx_tokens'] = num_ctx_tokens | ||
| self.iter_states['num_generation_tokens'] = num_generation_tokens | ||
| # Count the already-cached prefix for the sequences scheduled this iteration. | ||
| self.iter_states['cached_kv_tokens'] = sum(num_cached_tokens_per_seq) |
Collaborator
There was a problem hiding this comment.
Is it safe to remove this?
…puts on non-spec / non-mrope path PR NVIDIA#11943 added per-request Python LOAD_ATTR overhead in the generation-input prep path. PR NVIDIA#15633 added three unconditional _wait_for_model_engine_input_copy() syncs in PyExecutor's run loop. For dense non-multimodal / non-spec-decode workloads (e.g. gpt-oss-120b disagg dep2 at con=2048), the aggregate host-side ticks explain a ~6% throughput regression on GB200. Changes (Python-only, no runtime behavior change on the affected path): * Snapshot request.py_batch_idx and request.max_beam_num_tokens once per generation-request iteration in _prepare_tp_inputs; read sites now pay LOAD_FAST instead of LOAD_ATTR. The write-site `request.py_batch_idx = request.py_seq_slot` still fires after the reads, so snapshot-before-write ordering is preserved. * Skip get_runtime_tokens_per_gen_step() call when self.enable_spec_decode is False (callee is `lambda _: 1` when spec_config is None). * Gate 3 _wait_for_model_engine_input_copy() calls in PyExecutor on self._is_kv_manager_v2. Pre-NVIDIA#15633 the legacy V1 scheduler ran without any such sync; the sync is only needed for V2's host page-table. Verified on Lyris GB200 in the auto-perf-fix-pro workflow, 3-rep median: test: disagg-e2e-gb200_gpt-oss-120b-fp4_1k1k_con2048_ctx1_tp1_gen1_dep2_eplb0_mtp0_ccb-UCX good median (15d06c0): 105905.45 tok/s bad median (c25fa74): 99217.72 tok/s ToT median (a0c406f): 98526.45 tok/s Fix median (this diff): 103082.72 tok/s (+4.62% over ToT, +3.90% over bad) Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>
chenfeiz0326
force-pushed
the
repair-bot-bug6373447
branch
from
July 7, 2026 14:07
cd675a4 to
e7734fb
Compare
achartier
approved these changes
Jul 7, 2026
| self.runtime_draft_len) | ||
| else: | ||
| # Non-spec-decode path: callee is `lambda _: 1` when spec_config is None | ||
| # (see __init__:539). Inline the constant to skip the per-call frame setup. |
Collaborator
There was a problem hiding this comment.
Line number reference in comments go stale - replace by an actual description.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
PyTorchModelEngine._prepare_tgen_inputs/_prepare_inputsandPyExecutor: an extrasum(num_cached_tokens_per_seq)forcached_kv_tokens, an unconditionalget_runtime_tokens_per_gen_stepcall that resolves to a constant when spec decode is off, repeatedLOAD_ATTRon per-request Python attributes inside the generation loop, and an earlier-than-necessary_wait_for_model_engine_input_copysync on the legacy scheduler path. These accumulated CPU overhead per step, degrading Total_Token_Throughput vs 1.3.0rc17 on decode-heavy workloads.cached_kv_tokensbookkeeping, inlined the constantruntime_tokens_per_gen_step = 1when spec decode is disabled, snapshottedrequest.py_batch_idxandrequest.max_beam_num_tokensonce per request so read-sites useLOAD_FAST, and restored the pre-[None][feat] DSv4 follow-up: runtime KV and cache foundations #15633 behavior of deferring the input-copy wait for the legacy scheduler. These are minimal, semantics-preserving micro-optimizations that reverse the specific hot-path additions introduced since rc17 without changing scheduler or spec-decode behavior.Test plan
Links
Summary by CodeRabbit