[https://nvbugs/5440521][bug] Fix sequence slot allocation for attention DP - #6878
[https://nvbugs/5440521][bug] Fix sequence slot allocation for attention DP#6878Tabrizian wants to merge 1 commit into
Conversation
|
Caution Review failedThe pull request is closed. 📝 WalkthroughWalkthroughprepare_resources in seq_slot_manager.py now only allocates a new slot when Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant PyExecutor
participant SeqSlotManager
participant SlotAllocator
Client->>PyExecutor: submit request (llm_req)
PyExecutor->>SeqSlotManager: prepare_resources(llm_req)
SeqSlotManager->>SeqSlotManager: check llm_req.seq_slot\nand disagg state flags
alt seq_slot is None AND not (is_disagg_generation_init_state or is_disagg_generation_transmission_in_progress)
SeqSlotManager->>SlotAllocator: allocate_slot()
SlotAllocator-->>SeqSlotManager: slot
SeqSlotManager->>SeqSlotManager: set llm_req.seq_slot\nset llm_req.py_seq_slot
alt llm_req.return_perf_metrics
SeqSlotManager->>llm_req: set_first_scheduled_time()
end
else
SeqSlotManager->>SeqSlotManager: skip allocation
end
SeqSlotManager-->>PyExecutor: return
Estimated code review effort🎯 2 (Simple) | ⏱️ ~8 minutes Possibly related PRs
Suggested reviewers
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. 📜 Recent review detailsConfiguration used: .coderabbit.yaml 💡 Knowledge Base configuration:
You can enable these sources in your CodeRabbit configuration. 📒 Files selected for processing (1)
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
Status, 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: 1
🧹 Nitpick comments (2)
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py (2)
1-1: Prefer module-namespace imports for classes/enums.Guideline: maintain module namespace in imports. Consider importing the module and referencing symbols via the module to improve clarity and avoid name clashes.
Example:
from . import llm_request # type hints def get_needed_resource_to_completion(self, request: llm_request.LlmRequest) -> int: ... def free_resources(self, request: llm_request.LlmRequest) -> None: ... # usage if llm_req.seq_slot is None and llm_req.state != llm_request.LlmRequestState.DISAGG_GENERATION_INIT: ...Note: This change affects type hints and usage throughout the file; apply consistently if you adopt it.
19-21: Simplify condition: use != instead of “not … ==” for readability.No behavior change; just clearer intent.
Apply this diff:
- if llm_req.seq_slot is None and not llm_req.state == LlmRequestState.DISAGG_GENERATION_INIT: + if llm_req.seq_slot is None and llm_req.state != LlmRequestState.DISAGG_GENERATION_INIT:
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py(2 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.py: Python code must target Python 3.8+
Python indentation: 4 spaces, no tabs
Maintain module namespace in imports (from package.subpackage import foo; then use foo.SomeClass())
Python file names use snake_case
Python class names use PascalCase
Python functions/methods and local variables use snake_case; variables starting with a number get k_ prefix (e.g., k_99th_percentile)
Global variables use G_ prefixed UPPER_SNAKE_CASE (e.g., G_MY_GLOBAL)
Constants use UPPER_SNAKE_CASE in Python
Avoid shadowing variables from outer scopes in Python
Initialize all externally visible members of a Python class in init
Prefer docstrings for interfaces used outside a file; comments for local code
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes/variables inline with short docstrings
Avoid reflection when simple alternatives exist (e.g., prefer explicit parameters over dict(**locals()))
In try/except, catch the narrowest exceptions possible
For duck-typing with try/except, keep try body minimal and put logic in else
Files:
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py
🔇 Additional comments (1)
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py (1)
19-21: Seq-slot allocation looks correct — DISAGG_GENERATION_INIT should skip and will be allocated after transmission completesShort summary: I verified the execution paths — SeqSlotManager skips allocation only for DISAGG_GENERATION_INIT, and requests move to GENERATION_IN_PROGRESS (so get a slot) before the main resource allocation pass runs. No other disagg states need to be special-cased.
Relevant locations (checked):
- tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py — guard: if llm_req.seq_slot is None and not llm_req.state == LlmRequestState.DISAGG_GENERATION_INIT (allocates via SlotManager.add_slot).
- tensorrt_llm/_torch/pyexecutor/py_executor.py — main loop calls _prepare_disagg_gen_transmission_complete(...) before self.resource_manager.prepare_resources(scheduled_batch) (see the loop and ordering in the 660–760 and 1240–1336 ranges). _prepare_disagg_gen_transmission_complete transitions transmission-complete requests to LlmRequestState.GENERATION_IN_PROGRESS.
- tensorrt_llm/_torch/pyexecutor/py_executor.py — _prepare_disagg_gen_init(...) invokes resource managers for DISAGG_GENERATION_INIT requests (SeqSlotManager is called there but will intentionally skip allocation due to the guard).
- tensorrt_llm/_torch/pyexecutor/resource_manager.py — SlotManager.add_slot/remove_slot implement mapping and free_slots (add_slot used by SeqSlotManager; remove_slot called on free_resources).
Conclusion: the change is intentional and safe — we do not need to re-allocate earlier; allocation happens after the disagg transmission-complete → GENERATION_IN_PROGRESS transition. No additional states need to be skipped.
| @@ -1,4 +1,4 @@ | |||
| from .llm_request import LlmRequest | |||
| from .llm_request import LlmRequest, LlmRequestState | |||
There was a problem hiding this comment.
Add NVIDIA copyright header (current year).
Per coding guidelines, prepend the NVIDIA copyright header to all source files.
Apply this diff at the top of the file:
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
from .llm_request import LlmRequest, LlmRequestState📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| from .llm_request import LlmRequest, LlmRequestState | |
| # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | |
| from .llm_request import LlmRequest, LlmRequestState |
🤖 Prompt for AI Agents
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py lines 1-1: prepend the
standard NVIDIA copyright header with the current year to the top of this source
file; ensure the header is placed before any imports and matches project/legal
formatting (including year and company name), and save the file so the header
appears as the very first lines.
a863e7e to
2c1791e
Compare
|
/bot run |
|
PR_Github #15257 [ run ] triggered by Bot |
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py (1)
1-1: Add NVIDIA copyright header (current year).Per coding guidelines, prepend the NVIDIA copyright header to all source files.
Apply this diff at the top of the file:
+ # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + from .llm_request import LlmRequest, LlmRequestState
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py (1)
21-21: Prefer!=overnot ==for readability.Equivalent semantics; the inequality operator is clearer.
- if llm_req.seq_slot is None and not llm_req.state == LlmRequestState.DISAGG_GENERATION_INIT: + if llm_req.seq_slot is None and llm_req.state != LlmRequestState.DISAGG_GENERATION_INIT:
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py(2 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.py: Python code must target Python 3.8+
Python indentation: 4 spaces, no tabs
Maintain module namespace in imports (from package.subpackage import foo; then use foo.SomeClass())
Python file names use snake_case
Python class names use PascalCase
Python functions/methods and local variables use snake_case; variables starting with a number get k_ prefix (e.g., k_99th_percentile)
Global variables use G_ prefixed UPPER_SNAKE_CASE (e.g., G_MY_GLOBAL)
Constants use UPPER_SNAKE_CASE in Python
Avoid shadowing variables from outer scopes in Python
Initialize all externally visible members of a Python class in init
Prefer docstrings for interfaces used outside a file; comments for local code
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes/variables inline with short docstrings
Avoid reflection when simple alternatives exist (e.g., prefer explicit parameters over dict(**locals()))
In try/except, catch the narrowest exceptions possible
For duck-typing with try/except, keep try body minimal and put logic in else
Files:
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py
🧠 Learnings (1)
📚 Learning: 2025-08-12T10:28:57.320Z
Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-12T10:28:57.320Z
Learning: Applies to **/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py} : Prepend NVIDIA copyright header (current year) to all source files
Applied to files:
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py
🔇 Additional comments (3)
tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py (3)
1-1: LGTM: Importing LlmRequestState for state checks is correct.The added import is used appropriately in the allocation guard below.
19-21: Correctly skip slot allocation during DISAGG_GENERATION_INIT.This aligns allocation with post-init only, which is consistent with fixing attention DP init behavior.
19-21: SlotManager.remove_slot is safe for requests that never acquired a slot — no change required.remove_slot checks membership before popping and only mutates state when a slot exists, so calling free_resources for requests that skipped allocation (e.g., DISAGG_GENERATION_INIT) is a no-op.
- Evidence:
- tensorrt_llm/_torch/pyexecutor/resource_manager.py — SlotManager.remove_slot (conditional pop/add).
- tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py — prepare_resources skips DISAGG_GENERATION_INIT; free_resources calls remove_slot.
2c1791e to
9b47a9c
Compare
|
/bot run |
|
PR_Github #15350 [ run ] triggered by Bot |
|
PR_Github #15350 [ run ] completed with state |
|
I am quite new to trtllm, so I wonder whether we shall always keep logic alignment between py implementation and cpp implementation (e.g. also change the logic in |
|
@yifeizhang-c |
9b47a9c to
d56286b
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #15669 [ run ] triggered by Bot |
|
PR_Github #15669 [ run ] completed with state |
|
The current logic is similar to the changes used to fix
I am not sure whether any choice would make more sense compared with the other one. |
Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
d56286b to
52d285c
Compare
|
@yifeizhang-c I'm fine with that change too. Closing the PR in favour of your change. |
3 participants
Summary by CodeRabbit
Bug Fixes
Performance
Description
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.
Details
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--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-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-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.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline 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 in addition to running 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-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.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.