[nvbug/5322354] fix PD + MTP + overlap scheduler accuracy issue #6136
Conversation
WalkthroughThis update refines tensor initialization and offset management in speculative decoding within the model engine, ensuring explicit zeroing and selective copying for CUDA tensors. Additionally, it removes a redundant function call in the executor loop related to KV cache transmission completion, streamlining the control flow without altering public APIs. It also removes two test skip entries previously linked to a bug. Changes
Estimated code review effort2 (~15 minutes) Possibly related PRs
Suggested reviewers
Poem
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
1365-1391: LGTM: Excellent segmented request handling with clear documentation.The three-segment approach correctly addresses the accuracy issue by ensuring proper positioning of dummy requests at the end. The extensive comments provide valuable context for understanding the complex offset calculations.
Consider breaking the long comment lines to improve readability:
- # The order of request in Batch: ['requests that do not have previous batch', 'requests that already have previous batch', 'dummy requests'] + # The order of request in Batch: + # ['requests that do not have previous batch', 'requests that already have previous batch', 'dummy requests']Similar line breaking can be applied to other long comment lines (lines 1367-1368, 1370, 1373, 1375, 1377) to address static analysis warnings.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/model_engine.py(2 hunks)tensorrt_llm/_torch/pyexecutor/py_executor.py(0 hunks)tests/integration/test_lists/waives.txt(0 hunks)
💤 Files with no reviewable changes (2)
- tests/integration/test_lists/waives.txt
- tensorrt_llm/_torch/pyexecutor/py_executor.py
🧰 Additional context used
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/pyexecutor/model_engine.py
1365-1365: Line too long (156 > 120)
(E501)
1366-1366: Line too long (130 > 120)
(E501)
1367-1367: Line too long (159 > 120)
(E501)
1368-1368: Line too long (149 > 120)
(E501)
1370-1370: Line too long (128 > 120)
(E501)
1373-1373: Line too long (146 > 120)
(E501)
1375-1375: Line too long (163 > 120)
(E501)
1377-1377: Line too long (131 > 120)
(E501)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
1335-1337: LGTM: Critical fix for tensor initialization.The explicit zero-initialization of
previous_pos_id_offsets_cudaandprevious_kv_lens_offsets_cudaaddresses the root cause of the accuracy issue. This ensures a clean state before selective updates, preventing garbage values from affecting the overlap scheduler calculations.
1361-1364: LGTM: Clear variable naming and usage.The
new_tokens_lenvariable properly captures the tensor shape for subsequent offset calculations, improving code readability and maintainability.
|
/bot run |
|
PR_Github #12186 [ run ] triggered by Bot |
|
PR_Github #12186 [ run ] completed with state |
|
@yweng0828 Should the nvbug number be 5322354? |
yweng0828
changed the title
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
1318-1318: Remove the unused expression.The expression
len(request_ids)computes a value but doesn't assign it to anything, making it a no-op. This appears to be leftover code from a refactoring.- len(request_ids)
1363-1376: Excellent documentation but address line length violations.The detailed comments explaining the three-segment request organization are very helpful for understanding the complex overlap scheduling logic. However, several lines exceed the 120-character limit.
Consider reformatting the longer comment lines to comply with line length standards:
- # The order of requests in a batch: [context requests, generation requests] - # generation requests: ['requests that do not have previous batch', 'requests that already have previous batch', 'dummy requests'] - # 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving. - # 2) 'requests that already have previous batch': previous iteration's requests. - # 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. - # Therefore, both of self.previous_pos_id_offsets_cuda and self.previous_kv_lens_offsets_cuda are also 3 segments. - # For 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving. - # Set these requests' previous_pos_id_offsets and previous_kv_lens_offsets to '0' to skip the value changes in _preprocess_inputs. - # Already set to '0' during initialization. - # For 2) 'requests that already have previous batch': enable overlap scheduler. - # Set their previous_pos_id_offsets and previous_kv_lens_offsets according to new_tokens_lens_device and kv_len_offsets_device. - # For 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. - # Already set to '0' during initialization. + # The order of requests in a batch: [context requests, generation requests] + # generation requests: ['requests that do not have previous batch', + # 'requests that already have previous batch', 'dummy requests'] + # 1) 'requests that do not have previous batch': disable overlap scheduler or + # the first step in the generation server of disaggregated serving. + # 2) 'requests that already have previous batch': previous iteration's requests. + # 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. + # Therefore, both of self.previous_pos_id_offsets_cuda and + # self.previous_kv_lens_offsets_cuda are also 3 segments. + # For 1) 'requests that do not have previous batch': disable overlap scheduler or + # the first step in the generation server of disaggregated serving. + # Set these requests' previous_pos_id_offsets and previous_kv_lens_offsets + # to '0' to skip the value changes in _preprocess_inputs. + # Already set to '0' during initialization. + # For 2) 'requests that already have previous batch': enable overlap scheduler. + # Set their previous_pos_id_offsets and previous_kv_lens_offsets according + # to new_tokens_lens_device and kv_len_offsets_device. + # For 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. + # Already set to '0' during initialization.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py(3 hunks)
🧰 Additional context used
🧠 Learnings (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/pyexecutor/model_engine.py
1365-1365: Line too long (146 > 120)
(E501)
1366-1366: Line too long (161 > 120)
(E501)
1369-1369: Line too long (130 > 120)
(E501)
1370-1370: Line too long (165 > 120)
(E501)
1371-1371: Line too long (152 > 120)
(E501)
1374-1374: Line too long (149 > 120)
(E501)
🔇 Additional comments (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
1335-1337: LGTM: Explicit zero initialization improves reliability.The explicit initialization of CUDA tensors to zero at the start of the block is a good practice that ensures predictable behavior and eliminates potential bugs from relying on implicit state.
1377-1391: LGTM: Correct implementation of request segmentation logic.The tensor copying logic correctly implements the three-segment request organization described in the comments. The range calculations properly handle:
- Requests without previous batch (remain zero)
- Requests with previous batch (get offsets copied from device tensors)
- Dummy requests (remain zero)
This ensures accurate offset handling for the overlap scheduler in speculative decoding.
yweng0828
force-pushed
the
yweng/fix_pd_mtp2_overlap_scheduler_accuracy_issues
branch
from
f567db4 to
db15634
Compare
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
1311-1311: Remove the unused expression.The standalone expression
len(request_ids)doesn't serve any purpose and should be removed entirely if it's truly unused code.- len(request_ids)
1356-1384: Excellent segmentation logic with helpful comments.The detailed comments clearly explain the three-segment structure of requests and the rationale for zeroing vs. updating offset tensors. The indexing logic correctly isolates segment 2 (requests with previous batch) for offset updates while leaving segments 1 and 3 properly zeroed.
Consider breaking up the long lines to improve readability:
- # The order of requests in a batch: [context requests, generation requests] - # generation requests: ['requests that do not have previous batch', 'requests that already have previous batch', 'dummy requests'] - # 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving. - # 2) 'requests that already have previous batch': previous iteration's requests. - # 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. - # Therefore, both of self.previous_pos_id_offsets_cuda and self.previous_kv_lens_offsets_cuda are also 3 segments. - # For 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving. - # Set these requests' previous_pos_id_offsets and previous_kv_lens_offsets to '0' to skip the value changes in _preprocess_inputs. - # Already set to '0' during initialization. - # For 2) 'requests that already have previous batch': enable overlap scheduler. - # Set their previous_pos_id_offsets and previous_kv_lens_offsets according to new_tokens_lens_device and kv_len_offsets_device. - # For 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. - # Already set to '0' during initialization. + # The order of requests in a batch: [context requests, generation requests] + # generation requests: ['requests that do not have previous batch', + # 'requests that already have previous batch', 'dummy requests'] + # 1) 'requests that do not have previous batch': disable overlap scheduler + # or the first step in the generation server of disaggregated serving. + # 2) 'requests that already have previous batch': previous iteration's requests. + # 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. + # Therefore, both offset tensors are also segmented into 3 parts. + # For 1) 'requests that do not have previous batch': + # Set offset values to '0' to skip changes in _preprocess_inputs. + # Already set to '0' during initialization. + # For 2) 'requests that already have previous batch': enable overlap scheduler. + # Set offset values according to new_tokens_lens_device and kv_len_offsets_device. + # For 3) 'dummy requests': Already set to '0' during initialization.Also consider splitting the complex index calculations:
+ # Calculate start index for segment 2 (requests with previous batch) + segment2_start = (len(extend_requests) - len(extend_dummy_requests) - + previous_batch_len) * (1 + self.max_draft_len) + segment2_end = segment2_start + previous_batch_tokens + self.previous_pos_id_offsets_cuda[ - (len(extend_requests) - len(extend_dummy_requests) - - previous_batch_len) * (1 + self.max_draft_len): - (len(extend_requests) - len(extend_dummy_requests) - - previous_batch_len) * (1 + self.max_draft_len) + - previous_batch_tokens].copy_( + segment2_start:segment2_end].copy_( new_tokens_lens_device[self.previous_pos_indices_cuda[ non_blocking=True)
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/model_engine.py(3 hunks)tensorrt_llm/_torch/pyexecutor/py_executor.py(0 hunks)tests/integration/test_lists/waives.txt(0 hunks)
💤 Files with no reviewable changes (2)
- tests/integration/test_lists/waives.txt
- tensorrt_llm/_torch/pyexecutor/py_executor.py
🧰 Additional context used
🧠 Learnings (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/pyexecutor/model_engine.py
1358-1358: Line too long (146 > 120)
(E501)
1359-1359: Line too long (161 > 120)
(E501)
1362-1362: Line too long (130 > 120)
(E501)
1363-1363: Line too long (165 > 120)
(E501)
1364-1364: Line too long (152 > 120)
(E501)
1367-1367: Line too long (149 > 120)
(E501)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
1328-1330: Good fix for ensuring proper tensor initialization.The explicit zero-initialization of the offset tensors is crucial for the accuracy fix. This ensures these tensors start with known values rather than potentially containing stale data from previous iterations, which was likely causing the accuracy issues with the overlap scheduler.
|
PR_Github #12342 [ run ] triggered by Bot |
|
PR_Github #12342 [ run ] completed with state |
yweng0828
force-pushed
the
yweng/fix_pd_mtp2_overlap_scheduler_accuracy_issues
branch
from
db15634 to
856d94d
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
1311-1311: Remove the unused expression entirely.The line
len(request_ids)computes a value but doesn't use it. Since the variable assignment was removed, this expression serves no purpose and should be deleted completely.- len(request_ids)
1356-1385: Excellent documentation and logic improvement, with minor style issues.The detailed comments clearly explain the complex request segmentation and offset handling logic, which significantly improves code maintainability. The restructured offset calculations appear correct and should resolve the accuracy issues.
However, consider these improvements:
- Several lines exceed the 120-character limit (flagged by static analysis)
- Past review suggests using
num_tokensdirectly instead of the complex range calculation in lines 1374-1379- The variable name
num_extend_reqeust_wo_dummyhas a typo ("reqeust" → "request")Apply this diff to address the typo:
- num_extend_reqeust_wo_dummy = len(extend_requests) - len( + num_extend_request_wo_dummy = len(extend_requests) - len( extend_dummy_requests) self.previous_pos_id_offsets_cuda[ - (num_extend_reqeust_wo_dummy - previous_batch_len) * + (num_extend_request_wo_dummy - previous_batch_len) * (1 + self.max_draft_len):num_extend_reqeust_wo_dummy * + (1 + self.max_draft_len):num_extend_request_wo_dummy * (1 + self.max_draft_len)].copy_(Consider incorporating the past review suggestion to use
num_tokensdirectly as shown in the existing comment.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/model_engine.py(3 hunks)tensorrt_llm/_torch/pyexecutor/py_executor.py(0 hunks)tests/integration/test_lists/waives.txt(0 hunks)
💤 Files with no reviewable changes (2)
- tests/integration/test_lists/waives.txt
- tensorrt_llm/_torch/pyexecutor/py_executor.py
🧰 Additional context used
🧠 Learnings (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
Learnt from: amitz-nv
PR: #5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks is_adapter_in_cpu_cache() and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/pyexecutor/model_engine.py
1358-1358: Line too long (146 > 120)
(E501)
1359-1359: Line too long (161 > 120)
(E501)
1362-1362: Line too long (130 > 120)
(E501)
1363-1363: Line too long (165 > 120)
(E501)
1364-1364: Line too long (152 > 120)
(E501)
1367-1367: Line too long (149 > 120)
(E501)
🔇 Additional comments (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
1328-1330: Good improvement: explicit zero-initialization.Adding explicit zero-initialization of the offset tensors eliminates assumptions about their initial state and ensures correct behavior. This addresses the accuracy issue mentioned in the PR objectives.
|
/bot run |
|
PR_Github #12421 [ run ] triggered by Bot |
|
PR_Github #12421 [ run ] completed with state |
yweng0828
force-pushed
the
yweng/fix_pd_mtp2_overlap_scheduler_accuracy_issues
branch
from
856d94d to
3f9ac26
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
1322-1322: Remove dead code.Line 1322 computes
len(request_ids)but doesn't use the result. This appears to be leftover code that should be removed.- len(request_ids)
1382-1396: Logic looks correct, but fix the typo.The offset copying logic correctly implements the segmentation strategy described in the comments, ensuring only requests with previous batches get their offsets updated. However, there's a typo in the variable name.
- num_extend_reqeust_wo_dummy = len(extend_requests) - len( + num_extend_request_wo_dummy = len(extend_requests) - len( extend_dummy_requests) self.previous_pos_id_offsets_cuda[ - (num_extend_reqeust_wo_dummy - previous_batch_len) * + (num_extend_request_wo_dummy - previous_batch_len) * (1 + self.max_draft_len):num_extend_reqeust_wo_dummy * (1 + self.max_draft_len)].copy_( new_tokens_lens_device[self.previous_pos_indices_cuda[ non_blocking=True) self.previous_kv_lens_offsets_cuda[ - num_extend_reqeust_wo_dummy - + num_extend_request_wo_dummy - previous_batch_len:num_extend_reqeust_wo_dummy].copy_( kv_len_offsets_device[previous_slots], non_blocking=True)Also fix the remaining occurrence on line 1386:
- (1 + self.max_draft_len):num_extend_reqeust_wo_dummy * + (1 + self.max_draft_len):num_extend_request_wo_dummy *
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py(3 hunks)tensorrt_llm/_torch/pyexecutor/py_executor.py(0 hunks)
🧠 Learnings (2)
📓 Common learnings
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
Learnt from: amitz-nv
PR: #5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks is_adapter_in_cpu_cache() and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/pyexecutor/model_engine.py
1369-1369: Line too long (146 > 120)
(E501)
1370-1370: Line too long (161 > 120)
(E501)
1373-1373: Line too long (130 > 120)
(E501)
1374-1374: Line too long (165 > 120)
(E501)
1375-1375: Line too long (152 > 120)
(E501)
1378-1378: Line too long (149 > 120)
(E501)
💤 Files with no reviewable changes (1)
- tensorrt_llm/_torch/pyexecutor/py_executor.py
🧰 Additional context used
🧠 Learnings (2)
📓 Common learnings
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
Learnt from: amitz-nv
PR: #5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks is_adapter_in_cpu_cache() and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/pyexecutor/model_engine.py
1369-1369: Line too long (146 > 120)
(E501)
1370-1370: Line too long (161 > 120)
(E501)
1373-1373: Line too long (130 > 120)
(E501)
1374-1374: Line too long (165 > 120)
(E501)
1375-1375: Line too long (152 > 120)
(E501)
1378-1378: Line too long (149 > 120)
(E501)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
1339-1341: Good explicit initialization.Making the zero-initialization of the offset tensors explicit is a solid improvement that eliminates assumptions about their initial state. This should help prevent the accuracy issues mentioned in the PR.
1368-1381: Excellent documentation of complex logic.The detailed comments clearly explain the batch segmentation and offset handling strategy. This documentation is crucial for understanding why dummy requests must be placed at the end and how the offset tensors are managed for different request types.
|
PR_Github #12526 [ run ] triggered by Bot |
|
PR_Github #12526 [ run ] completed with state |
|
/bot run |
|
PR_Github #12558 [ run ] triggered by Bot |
|
PR_Github #12558 [ run ] completed with state |
yweng0828
force-pushed
the
yweng/fix_pd_mtp2_overlap_scheduler_accuracy_issues
branch
from
19e8b3c to
6f750da
Compare
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
1367-1381: Excellent documentation, but consider line lengthThe detailed comments explaining the request segmentation and offset handling logic are very valuable for code maintainability. This level of documentation is crucial for understanding the complex overlap scheduling behavior.
Consider wrapping the longer comment lines to stay within the 120-character limit:
- # The order of requests in a batch: [context requests, generation requests] - # generation requests: ['requests that do not have previous batch', 'requests that already have previous batch', 'dummy requests'] - # 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving. - # 2) 'requests that already have previous batch': previous iteration's requests. - # 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. - # Therefore, both of self.previous_pos_id_offsets_cuda and self.previous_kv_lens_offsets_cuda are also 3 segments. - # For 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving. - # Set these requests' previous_pos_id_offsets and previous_kv_lens_offsets to '0' to skip the value changes in _preprocess_inputs. - # Already set to '0' during initialization. - # For 2) 'requests that already have previous batch': enable overlap scheduler. - # Set their previous_pos_id_offsets and previous_kv_lens_offsets according to new_tokens_lens_device and kv_len_offsets_device. - # For 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. - # Already set to '0' during initialization. + # The order of requests in a batch: [context requests, generation requests] + # generation requests: ['requests that do not have previous batch', + # 'requests that already have previous batch', 'dummy requests'] + # 1) 'requests that do not have previous batch': disable overlap scheduler or + # the first step in the generation server of disaggregated serving. + # 2) 'requests that already have previous batch': previous iteration's requests. + # 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. + # Therefore, both previous_pos_id_offsets_cuda and previous_kv_lens_offsets_cuda + # are also 3 segments. + # For 1) 'requests that do not have previous batch': disable overlap scheduler or + # the first step in the generation server of disaggregated serving. + # Set these requests' previous_pos_id_offsets and previous_kv_lens_offsets + # to '0' to skip the value changes in _preprocess_inputs. + # Already set to '0' during initialization. + # For 2) 'requests that already have previous batch': enable overlap scheduler. + # Set their previous_pos_id_offsets and previous_kv_lens_offsets according + # to new_tokens_lens_device and kv_len_offsets_device. + # For 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. + # Already set to '0' during initialization.
1382-1397: Fix variable name typo, otherwise logic looks correctThe offset copying logic correctly implements the segmented approach described in the comments, properly targeting requests with previous batches while leaving other segments zeroed. This addresses the core accuracy issue mentioned in the PR.
There's a typo in the variable name:
- num_extend_reqeust_wo_dummy = len(extend_requests) - len( + num_extend_request_wo_dummy = len(extend_requests) - len( extend_dummy_requests) self.previous_pos_id_offsets_cuda[ - (num_extend_reqeust_wo_dummy - previous_batch_len) * + (num_extend_request_wo_dummy - previous_batch_len) * - (1 + self.max_draft_len):num_extend_reqeust_wo_dummy * + (1 + self.max_draft_len):num_extend_request_wo_dummy * (1 + self.max_draft_len)].copy_( new_tokens_lens_device[self.previous_pos_indices_cuda[ non_blocking=True) self.previous_kv_lens_offsets_cuda[ - num_extend_reqeust_wo_dummy - + num_extend_request_wo_dummy - - previous_batch_len:num_extend_reqeust_wo_dummy].copy_( + previous_batch_len:num_extend_request_wo_dummy].copy_( kv_len_offsets_device[previous_slots], non_blocking=True)
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/model_engine.py(2 hunks)tensorrt_llm/_torch/pyexecutor/py_executor.py(0 hunks)tests/integration/test_lists/waives.txt(0 hunks)
🧠 Learnings (2)
📓 Common learnings
Learnt from: yechank-nvidia
PR: NVIDIA/TensorRT-LLM#6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.703Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
Learnt from: yechank-nvidia
PR: #6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.703Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using from_shared_tensor() is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call strip_for_generation() to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
Learnt from: amitz-nv
PR: #5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks is_adapter_in_cpu_cache() and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/pyexecutor/model_engine.py
1369-1369: Line too long (146 > 120)
(E501)
1370-1370: Line too long (161 > 120)
(E501)
1373-1373: Line too long (130 > 120)
(E501)
1374-1374: Line too long (165 > 120)
(E501)
1375-1375: Line too long (152 > 120)
(E501)
1378-1378: Line too long (149 > 120)
(E501)
💤 Files with no reviewable changes (2)
- tensorrt_llm/_torch/pyexecutor/py_executor.py
- tests/integration/test_lists/waives.txt
🧰 Additional context used
🧠 Learnings (2)
📓 Common learnings
Learnt from: yechank-nvidia
PR: NVIDIA/TensorRT-LLM#6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.703Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
Learnt from: yechank-nvidia
PR: #6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.703Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using from_shared_tensor() is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call strip_for_generation() to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
Learnt from: amitz-nv
PR: #5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks is_adapter_in_cpu_cache() and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/pyexecutor/model_engine.py
1369-1369: Line too long (146 > 120)
(E501)
1370-1370: Line too long (161 > 120)
(E501)
1373-1373: Line too long (130 > 120)
(E501)
1374-1374: Line too long (165 > 120)
(E501)
1375-1375: Line too long (152 > 120)
(E501)
1378-1378: Line too long (149 > 120)
(E501)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
1339-1341: LGTM: Explicit tensor initialization improves reliabilityThe explicit zero-initialization of offset tensors addresses potential undefined behavior and ensures a known starting state for overlap scheduling. This change aligns well with the PR's objective to fix accuracy issues in the PD + MTP + overlap scheduler combination.
|
PR_Github #12580 [ run ] triggered by Bot |
|
PR_Github #12580 [ run ] completed with state |
|
/bot run |
|
PR_Github #12583 [ run ] triggered by Bot |
|
PR_Github #12583 [ run ] completed with state |
Signed-off-by: Yue Weng <[email protected]>
Signed-off-by: Yue Weng <[email protected]>
Signed-off-by: Yue Weng <[email protected]>
Signed-off-by: Yue Weng <[email protected]>
Signed-off-by: Yue Weng <[email protected]>
yweng0828
force-pushed
the
yweng/fix_pd_mtp2_overlap_scheduler_accuracy_issues
branch
from
6f750da to
2d1e020
Compare
|
/bot run |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
371-383: Excellent documentation, minor formatting suggestion.The detailed comments clearly explain the complex three-segment request structure, which is crucial for understanding this overlap scheduler logic. This documentation will help prevent similar accuracy issues in the future.
Consider breaking some longer comment lines to adhere to the 120-character limit:
- # The order of requests in a batch: [context requests, generation requests] - # generation requests: ['requests that do not have previous batch', 'requests that already have previous batch', 'dummy requests'] + # The order of requests in a batch: [context requests, generation requests] + # generation requests: [ + # 'requests that do not have previous batch', + # 'requests that already have previous batch', + # 'dummy requests' + # ]
385-399: Core fix looks correct, but please address the typo.The tensor update logic correctly implements the three-segment structure to fix the accuracy issue. The range calculations ensure that dummy requests are properly placed at the end and only requests with previous batch get their offsets updated.
Fix the typo in the variable name:
- num_extend_reqeust_wo_dummy = len(extend_requests) - len( + num_extend_request_wo_dummy = len(extend_requests) - len( extend_dummy_requests) self.previous_pos_id_offsets_cuda[ - (num_extend_reqeust_wo_dummy - previous_batch_len) * + (num_extend_request_wo_dummy - previous_batch_len) * (1 + self.max_draft_len):num_extend_reqeust_wo_dummy * + (1 + self.max_draft_len):num_extend_request_wo_dummy * (1 + self.max_draft_len)].copy_(And update the subsequent usage:
self.previous_kv_lens_offsets_cuda[ - num_extend_reqeust_wo_dummy - + num_extend_request_wo_dummy - previous_batch_len:num_extend_reqeust_wo_dummy].copy_( + previous_batch_len:num_extend_request_wo_dummy].copy_(
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/model_engine.py(2 hunks)tensorrt_llm/_torch/pyexecutor/py_executor.py(0 hunks)tests/integration/test_lists/waives.txt(0 hunks)
🧠 Learnings (2)
📓 Common learnings
Learnt from: yechank-nvidia
PR: NVIDIA/TensorRT-LLM#6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.703Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
Learnt from: yechank-nvidia
PR: #6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.703Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using from_shared_tensor() is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call strip_for_generation() to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
Learnt from: amitz-nv
PR: #5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks is_adapter_in_cpu_cache() and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/pyexecutor/model_engine.py
1372-1372: Line too long (146 > 120)
(E501)
1373-1373: Line too long (161 > 120)
(E501)
1376-1376: Line too long (130 > 120)
(E501)
1377-1377: Line too long (165 > 120)
(E501)
1378-1378: Line too long (152 > 120)
(E501)
1381-1381: Line too long (149 > 120)
(E501)
💤 Files with no reviewable changes (2)
- tensorrt_llm/_torch/pyexecutor/py_executor.py
- tests/integration/test_lists/waives.txt
🧰 Additional context used
🧠 Learnings (2)
📓 Common learnings
Learnt from: yechank-nvidia
PR: NVIDIA/TensorRT-LLM#6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.703Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
tensorrt_llm/_torch/pyexecutor/model_engine.py (2)
Learnt from: yechank-nvidia
PR: #6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.703Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using from_shared_tensor() is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call strip_for_generation() to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
Learnt from: amitz-nv
PR: #5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.374Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks is_adapter_in_cpu_cache() and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/pyexecutor/model_engine.py
1372-1372: Line too long (146 > 120)
(E501)
1373-1373: Line too long (161 > 120)
(E501)
1376-1376: Line too long (130 > 120)
(E501)
1377-1377: Line too long (165 > 120)
(E501)
1378-1378: Line too long (152 > 120)
(E501)
1381-1381: Line too long (149 > 120)
(E501)
🔇 Additional comments (1)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
342-344: LGTM: Explicit tensor initialization fixes accuracy issue.The explicit zero initialization of
previous_pos_id_offsets_cudaandprevious_kv_lens_offsets_cudaaddresses the core accuracy problem mentioned in the PR objectives. This ensures proper segmentation handling for requests without previous batch, with previous batch, and dummy requests.
|
PR_Github #12636 [ run ] triggered by Bot |
|
PR_Github #12636 [ run ] completed with state |
|
Resolving conversations after confirming with @yweng0828 offline, thanks. |
…IA#6136) Signed-off-by: Yue Weng <[email protected]> Signed-off-by: Shreyas Misra <[email protected]>
…IA#6136) Signed-off-by: Yue Weng <[email protected]> Signed-off-by: Ransiki Zhang <[email protected]>
…IA#6136) Signed-off-by: Yue Weng <[email protected]> Signed-off-by: Lanyu Liao <[email protected]>
Description
Fixed the accuracy issue of PD + MTP + overlap scheduler.
The root cause is that the update of
previous_pos_id_offsets_cudaandprevious_kv_lens_offsets_cudais wrong. The dummy request should be placed at the end.Test Coverage
GitHub Bot Help
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
run [--disable-fail-fast --skip-test --stage-list "A10-1, xxx" --gpu-type "A30, H100_PCIe" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-[Post-Merge]-1, xxx"]Launch build/test pipelines. All previously running jobs will be killed.
--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests. Will also run L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-[Post-Merge]-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-[Post-Merge]-1, xxx".For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.md.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.
Summary by CodeRabbit