[https://nvbugs/5394392][fix] Enlarge scheduler capacity under disagg bs == 1 #6537
Conversation
📝 WalkthroughWalkthroughAdds conditional +1 to max_num_sequences in PyTorch executor/sampler when gen batch size is 1 with KV cache present, skips DISAGG_GENERATION_INIT requests during sequence-slot assignment in C++ scheduler, and introduces a new integration test and test-list entry for gen batch size 1 (TinyLlama). Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Scheduler
participant AssignReqSeqSlots
Client->>Scheduler: Submit requests
Scheduler->>AssignReqSeqSlots: Assign sequence slots
loop For each request
AssignReqSeqSlots->>AssignReqSeqSlots: Check state
alt state == DISAGG_GENERATION_INIT
AssignReqSeqSlots-->>AssignReqSeqSlots: continue (skip assignment)
else
AssignReqSeqSlots->>AssignReqSeqSlots: Evaluate isReqNew / acquire slot
end
end
AssignReqSeqSlots-->>Scheduler: Assigned slots (excluding INIT)
sequenceDiagram
participant Caller
participant PyExecUtil
Caller->>PyExecUtil: create_py_executor_instance(max_num_sequences, kv_cache_mgr)
alt max_num_sequences==1 and kv_cache_mgr present
PyExecUtil->>PyExecUtil: max_num_sequences += 1
end
PyExecUtil-->>Caller: Executor instance
Caller->>PyExecUtil: create_torch_sampler_args(max_num_sequences, kv_cache_cfg)
alt max_num_sequences==1 and kv_cache_cfg present
PyExecUtil->>PyExecUtil: max_num_sequences += 1
end
PyExecUtil-->>Caller: Sampler args
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~15 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. ✨ 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: 0
♻️ Duplicate comments (1)
tensorrt_llm/_torch/pyexecutor/_util.py (1)
565-568: Same adjustment duplicated here – apply the helper for consistencyThis block repeats the logic discussed above. Once a helper like
_effective_max_num_sequences()is introduced, replace the open-coded calculation here to guarantee both the sampler and the slot/scheduler stay in lock-step.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/_util.py (1)
509-513: Factor out the “extra-capacity” formula to avoid duplication and future driftThe same
if executor_config.max_batch_size == 1: max_num_sequences += mapping.pp_sizepattern now lives in two different code paths (create_py_executor_instanceandcreate_torch_sampler_args). If this heuristic ever needs tuning (e.g. the required headroom changes again), the risk of the two locations diverging is high.Consider extracting a small helper, e.g.
def _effective_max_num_sequences(max_batch_size: int, pp_size: int) -> int: cap = max_batch_size * pp_size if max_batch_size == 1: cap += pp_size return capand reuse it in both call sites.
This keeps the policy in one place, makes the intent explicit, and simplifies unit testing for edge cases (pp_size > 1,max_batch_size > 1, etc.).
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tensorrt_llm/_torch/pyexecutor/_util.py(2 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: in tensorrt_llm/executor/worker.py, the lora adapter cache optimization logic that checks `is_adapte...
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.402Z
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.
Applied to files:
tensorrt_llm/_torch/pyexecutor/_util.py
yifeizhang-c
force-pushed
the
dev-yifeiz-fix-bs1-nonstop
branch
3 times, most recently
from
ca73d42 to
69a5319
Compare
yifeizhang-c
changed the title
yifeizhang-c
changed the title
yifeizhang-c
changed the title
coderabbitai
bot
requested review from
mikeiovine,
pcastonguay,
wili-65535 and
yuxianq
yifeizhang-c
force-pushed
the
dev-yifeiz-fix-bs1-nonstop
branch
3 times, most recently
from
be42893 to
4f68507
Compare
|
/bot run |
|
PR_Github #14084 [ run ] triggered by Bot |
|
PR_Github #14084 [ run ] completed with state |
yifeizhang-c
force-pushed
the
dev-yifeiz-fix-bs1-nonstop
branch
from
4f68507 to
00f167f
Compare
|
/bot run |
|
PR_Github #14096 [ run ] triggered by Bot |
|
PR_Github #14096 [ run ] completed with state |
yifeizhang-c
force-pushed
the
dev-yifeiz-fix-bs1-nonstop
branch
from
00f167f to
f8cc579
Compare
|
/bot run |
|
PR_Github #14131 [ run ] triggered by Bot |
|
@Shixiaowei02 @qiaoxj07 Can you also help take a look at the current PR? |
|
PR_Github #14131 [ run ] completed with state |
|
/bot run |
|
PR_Github #15051 [ run ] triggered by Bot |
|
PR_Github #15051 [ run ] completed with state |
|
@yifeizhang-c Do you mind responding to the question here? #6537 (comment) Increasing |
Increasing Sorry for not replying the previous comment in time, as the suggestion indeed will make the code changes more concise, and I was analyzing into it. To resolve the disagg bs==1 non-stop issue, increasing the |
Signed-off-by: Yifei Zhang <[email protected]>
|
Updated the logic to move |
|
/bot run --disable-fail-fast |
|
PR_Github #15326 [ run ] triggered by Bot |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
cpp/tensorrt_llm/batch_manager/assignReqSeqSlots.cpp (1)
35-36: Nit: punctuate the comment and optionally add a trace log to aid debugging.A terminal period aligns with our style. The trace log makes it easier to confirm that INIT requests are indeed skipped at runtime.
Apply:
- // Skip assigning sequence slot for DISAGG_GENERATION_INIT request - continue; + // Skip assigning sequence slot for DISAGG_GENERATION_INIT request. + TLLM_LOG_TRACE("Skipping seq-slot assignment for DISAGG_GENERATION_INIT request ID %lu", llmReq->mRequestId); + continue;
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these settings in your CodeRabbit configuration.
📒 Files selected for processing (1)
cpp/tensorrt_llm/batch_manager/assignReqSeqSlots.cpp(1 hunks)
🧰 Additional context used
📓 Path-based instructions (4)
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh}: In C++, close namespaces with a comment naming the namespace (e.g., } // namespace foo)
Prefer const/constexpr variables over #define for constants
Declare variables const if not modified after initialization
Use Allman brace style in C++
C++ filenames use lowerCamelCase and must be case-insensitively unique within a build target
C++ type names use UpperCamelCase
Local variables, methods, and namespaces use lowerCamelCase
Global non-static variables not in anonymous namespace use gPrefix lowerCamelCase (e.g., gExample)
Static globals or globals in anonymous namespaces use sPrefix lowerCamelCase
Locally visible static variables start with 's' (e.g., static std::once_flag sFlag;)
Member variables use mPrefix lowerCamelCase; public members may omit but are encouraged to use 'm'
Constants (enums, global/static/function-scope magic numbers) use kPREFIXED_UPPER_SNAKE (e.g., kDIGIT_NUM)
If macros are unavoidable, use UPPER_SNAKE_CASE (prefer constants over #define)
Constructor parameter that conflicts with a public member name gets trailing underscore (foo_)
Literal suffixes should be uppercase (e.g., 1234L not 1234l)
C++: use spaces only; indent 4 spaces
Run clang-format (LLVM style) before submitting; wrap lines at 120 characters
If formatting must be bypassed, use // clang-format off/on around the section
Prefer smart pointers; use unique_ptr for sole ownership, shared_ptr for shared; weak_ptr only in exceptional cases
Do not use deprecated pre-C++11 smart pointers
Use C++ style comments; avoid C comments except special inline cases; prefer // single-line
Capitalize and punctuate full-sentence comments
Follow Doxygen rules: use //! for comments and //!< for members in C++
Disable code with #if/#endif and mnemonic conditions; avoid commented-out code; avoid dead code
Do not throw exceptions across library boundaries
Use least-forceful casts; avoid removing const/volatile; avoid C-style and functional casts (except constructors); p...
Files:
cpp/tensorrt_llm/batch_manager/assignReqSeqSlots.cpp
**/*.{cpp,cxx,cc,cu}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.{cpp,cxx,cc,cu}: Avoid literal values except for 0, nullptr, true, false; use named constexpr for other literals
Place semicolon of empty for/while loop on a new line
Always use brace-delimited bodies for switch/while/do-for/if/else
Use inline C comments in argument lists when parameter meaning is unclear (e.g., /* checkForErrors = */ false)
Do not use assignment in subexpressions (e.g., if (x = y) ... is forbidden)
Switch on enums should enumerate all values and omit default to catch new values at compile time
Structure switch statements; prohibit fallthrough except between empty cases; each case ends with break or throw; return at end of case not allowed; put break inside braces for compound case
Prefer anonymous namespaces over static for internal linkage of functions
Every defined function must be called at least once (no unused methods)
Files:
cpp/tensorrt_llm/batch_manager/assignReqSeqSlots.cpp
**/*.{h,hpp,hxx,hh,cuh,cpp,cxx,cc,cu}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Parameter names must be consistent between declarations and definitions
Files:
cpp/tensorrt_llm/batch_manager/assignReqSeqSlots.cpp
**/*.{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:
cpp/tensorrt_llm/batch_manager/assignReqSeqSlots.cpp
⏰ 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)
cpp/tensorrt_llm/batch_manager/assignReqSeqSlots.cpp (1)
33-37: Skipping DISAGG_GENERATION_INIT from seq-slot assignment — verified safeShort summary: I inspected the call flow and usages of mSeqSlot / perf timestamps. Skipping seq-slot assignment for DISAGG_GENERATION_INIT is intentional and safe — no downstream code path reads mSeqSlot or relies on firstScheduledTime while a request is in the INIT state; slots are assigned later when transmission completes.
Relevant findings (key locations inspected)
- cpp/tensorrt_llm/batch_manager/assignReqSeqSlots.cpp (lines ~29–46): the continue for isDisaggGenerationInitState() is present; setFirstScheduledTime() is only called for isReqNew (context init or isDisaggGenerationTransmissionComplete).
- cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp:
- capacity + micro-batching flow: capacity scheduler returns fittingRequests and fittingDisaggGenInitRequests; prepareDisaggGenInitRequests(...) handles KV transfer for DISAGG_GENERATION_INIT (no seq-slot use) and mAssignReqSeqSlots is called only for microBatchScheduler output (currRequests) — see the calls around lines ~1035 and ~1064, and prepareDisaggGenInitRequests around ~1520–1620.
- cpp/tensorrt_llm/batch_manager/sequenceSlotManager.cpp: getSequenceSlot(startFlag, sequenceId) only assigns on startFlag; freeSequenceSlot / freeIdleSequenceSlots manage lifetimes.
- Code paths that dereference mSeqSlot (e.g., createNewDecoderRequests.cpp, makeDecodingBatchInputOutput.cpp, runtime/gptDecoderBatched.cpp, runtime/decoderState.cpp, transformerBuffers.cpp, handleGenerationLogits.cpp, runtimeBuffers.cpp) operate on requests produced by the microBatchScheduler (scheduled context/generation requests) — i.e., requests that should have a slot by the time those functions run.
- Perf metrics: setFirstScheduledTime is invoked only when a slot is newly assigned (isReqNew && getReturnPerfMetrics()); firstScheduledTime is used for reporting/serialization (executor serialization, triton reporting), not for control flow that would break if absent during INIT.
- Python-side mirror: tensorrt_llm/_torch/pyexecutor/seq_slot_manager.py also explicitly skips DISAGG_GENERATION_INIT, and assigns slot + first-scheduled-time later (consistent behavior).
- Integration test & scheduler guard: the repo contains test_disaggregated_genbs1 and the Python-side scheduler_capacity bump for bs==1 (tensorrt_llm/_torch/pyexecutor/_util.py) to prevent scheduler deadlocks — complementary to this change.
Conclusion / action
- No code changes required here; the continue is correct and safe. Approve this change.
|
Thanks @yifeizhang-c, I think the pytorch runtime doesn't use the C++ batch scheduler (feel free to correct me). I think the reason that In this case, it would just schedule the dummy request and never reach the generation init request since there is a dummy request added in every iteration. Here is a draft of the scheduler changes which I think is needed for this: https://github.com/NVIDIA/TensorRT-LLM/compare/main...Tabrizian:TensorRT-LLM:user/imant/schedulerChange?expand=1 Please let me know if you have any feedback. |
|
PR_Github #15326 [ run ] completed with state |
Currently
I am not clear about the logic behind the implementation of |
yifeizhang-c
changed the title
|
/bot run |
|
PR_Github #15394 [ run ] triggered by Bot |
|
PR_Github #15394 [ run ] completed with state |
… bs == 1 (NVIDIA#6537) Signed-off-by: Yifei Zhang <[email protected]>
… bs == 1 (NVIDIA#6537) Signed-off-by: Yifei Zhang <[email protected]> Signed-off-by: Wangshanshan <[email protected]>
… bs == 1 (NVIDIA#6537) Signed-off-by: Yifei Zhang <[email protected]> Signed-off-by: Wangshanshan <[email protected]>
… bs == 1 (NVIDIA#6537) Signed-off-by: Yifei Zhang <[email protected]>
… bs == 1 (NVIDIA#6537) Signed-off-by: Yifei Zhang <[email protected]> Signed-off-by: Wangshanshan <[email protected]>
… bs == 1 (NVIDIA#6537) Signed-off-by: Yifei Zhang <[email protected]>
… bs == 1 (NVIDIA#6537) Signed-off-by: Yifei Zhang <[email protected]>
… bs == 1 (NVIDIA#6537) Signed-off-by: Yifei Zhang <[email protected]> Signed-off-by: Wangshanshan <[email protected]>
Description
Under disagg, if generation server runs with
bs == 1, the dummyGENERATION_IN_PROGRESSrequest for attention dp will preventDISAGG_GENERATION_INITfrom being scheduled, thus letting generation server running in an endless cycle. This PR enlarges scheduler capacity and related resources to be with at least capacity == 2.Besides, originally
py_executorlogic assign newSEQ_SLOTresource onDISAGG_GEN_INITstate. This PR delays the assignment toDISAGG_TRANS_COMPLETEstate.Summary by CodeRabbit
Bug Fixes
Tests