[None][fix] Add TestServePrefixAwareScheduling base on LMBenchmark/synthetic-multi-round-qa - #13578
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #45993 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis change set introduces prefix-reuse summary caching to the scheduler to avoid redundant KV-cache radix-tree walks, optimizes a C++ trie node-clearing operation to prevent stale cascade-pruning, and adds comprehensive end-to-end and unit tests for prefix-aware scheduling behavior under various configurations. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Scheduler as PyCapacityScheduler
participant SummaryCache as Summary Cache<br/>(by request)
participant BenefitChk as _beneficial_to_skip()
participant BlkMgr as Block Manager
participant KVMgr as KV Manager
Client->>Scheduler: schedule(requests)
Scheduler->>SummaryCache: create caches for policy run
loop For each request
Scheduler->>BenefitChk: call with request
alt Summary not cached
BenefitChk->>KVMgr: analyze_prefix_reuse (normal)
KVMgr-->>BenefitChk: PrefixReuseSummary
BenefitChk->>SummaryCache: store summary
else Summary cached
BenefitChk->>SummaryCache: retrieve cached summary
SummaryCache-->>BenefitChk: PrefixReuseSummary
end
BenefitChk-->>Scheduler: skip decision
alt Should schedule
Scheduler->>BlkMgr: enough_available_blocks(req, cached_summary)
BlkMgr->>SummaryCache: lookup summary
SummaryCache-->>BlkMgr: PrefixReuseSummary
BlkMgr-->>Scheduler: boolean result
Scheduler->>BlkMgr: prepare_blocks_if_schedulable(req, cached_summary)
BlkMgr->>SummaryCache: lookup summary
SummaryCache-->>BlkMgr: PrefixReuseSummary
BlkMgr-->>Scheduler: block allocation dict
Scheduler->>BlkMgr: decrement_reserved_blocks(req, cached_summary)
BlkMgr->>SummaryCache: lookup summary
SummaryCache-->>BlkMgr: PrefixReuseSummary
end
end
Scheduler-->>Client: scheduled requests
Estimated Code Review Effort🎯 4 (Complex) | ⏱️ ~50 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 3❌ Failed checks (2 warnings, 1 inconclusive)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py (1)
622-800:⚠️ Potential issue | 🟡 MinorAdd this test to the QA scheduled test list per coding guidelines.
Per
tests/integration/test_lists/qa/README.md, "all new test cases should be added" tollm_function_core.txtfor single-node multi-GPU functional scenarios. The newTestServePrefixAwareSchedulingclass is a functional E2E test that validates trtllm-serve behavior and should be included in QA scheduled runs. The test-db wiring inl0_h100.ymlandl0_b200.ymlis correct, but entries must also be added totests/integration/test_lists/qa/llm_function_core.txt(or other appropriate QA list if multi-node coverage is desired).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py` around lines 622 - 800, The new TestServePrefixAwareScheduling E2E test class is not listed in the QA scheduled test list; add an entry for this test to tests/integration/test_lists/qa/llm_function_core.txt (or another appropriate QA list for multi-node if intended) so it runs in scheduled QA; specifically include the test class name TestServePrefixAwareScheduling (or its fully qualified test path if project uses module-style entries) in the llm_function_core.txt file following the formatting of existing entries.
🧹 Nitpick comments (1)
tests/integration/test_lists/test-db/l0_h100.yml (1)
170-178: Decide whether one representative case belongs in scheduled QA.These entries cover pre-merge, but the diff does not mirror the new prefix-aware test into
tests/integration/test_lists/qa/. If this is intended as ongoing release coverage for core serving behavior, add one representative case there; otherwise QA list updates are optional.As per coding guidelines, "If the change adds or materially alters an integration test under tests/integration/defs/ (or otherwise affects what QA should run on a schedule), call out whether an entry is needed under tests/integration/test_lists/qa/."
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/integration/test_lists/test-db/l0_h100.yml` around lines 170 - 178, Decide whether this new prefix-aware scheduling E2E test should be run on the scheduled QA list; if it is intended as ongoing release coverage for core serving behavior, add one representative case from kv_cache/test_prefix_aware_scheduling.py (for example TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix or one specific parametrized variant like [guaranteed-chunked]) into tests/integration/test_lists/qa/ so scheduled QA mirrors the pre-merge list; if it is not intended for scheduled QA, leave the QA list unchanged and document that the tests are pre-merge only.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py`:
- Around line 824-832: The code currently charges the full chunk slack
(remaining_space) to num_ctx_tokens even when the request has fewer draft
tokens, over-consuming capacity; fix by charging only the number of draft tokens
actually kept: compute kept_drafts = min(req.num_draft_tokens, remaining_space)
(ensure kept_drafts >= 0), then do num_ctx_tokens += kept_drafts and set
draft_discard = req.num_draft_tokens - kept_drafts (clamped to >=0). Update the
block that references remaining_space, num_ctx_tokens, req.num_draft_tokens, and
draft_discard to use kept_drafts instead of adding remaining_space directly.
In `@tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py`:
- Around line 514-515: The test helper _parse_csv_metrics() currently returns
None when all parsed TTFTs (ttfts) are invalid, which causes
_run_and_assert_stage() to report "no completed requests" incorrectly; update
_parse_csv_metrics() so that if ttfts is empty or all values are
NaN/Inf/unparseable it returns a metrics dict (not None) containing at least
num_requests (total CSV rows processed) and nan_count (count of invalid TTFT
rows) along with any other metrics, so failures point to corrupted latency
output; make the same change in the other identical block handling ttfts (the
second occurrence referenced in the comment) to ensure both code paths return a
metrics dict rather than None.
- Around line 424-431: The try/except around the health poll is too broad and
masks unrelated errors; replace the bare "except Exception:" that wraps the
req_lib.get(...) call with "except RequestException:" (the tests already import
RequestException from requests.exceptions) so only request-layer failures are
swallowed and other bugs in the health-check/watchdog path (e.g., failures
setting last_health_ok) will surface; keep the existing timeout and the logic
that sets last_health_ok unchanged.
- Around line 308-312: The benchmark harness must consistently resolve and run
the installed LMBenchmark using the virtualenv created by the llm_venv fixture:
change _resolve_lmbenchmark_script/fixture behavior so it returns both the
resolved LMBENCHMARK_SCRIPT path and the venv Python interpreter (or have the
fixture launch the benchmark directly), and ensure subprocess.Popen is invoked
with that returned venv interpreter instead of sys.executable (references:
llm_venv, _resolve_lmbenchmark_script, LMBENCHMARK_SCRIPT, script_dir, cmd,
subprocess.Popen). Also tighten exception handling: replace the broad "except
Exception:" around the /health check with "except RequestException:"
(RequestException is already imported) and narrow the watchdog loop catch from
Exception to specific subprocess-related exceptions such as
subprocess.SubprocessError, OSError, and TimeoutError to avoid masking
pytest/test-control errors (references: the /health check block and the watchdog
loop that manages proc.poll/communicate).
---
Outside diff comments:
In `@tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py`:
- Around line 622-800: The new TestServePrefixAwareScheduling E2E test class is
not listed in the QA scheduled test list; add an entry for this test to
tests/integration/test_lists/qa/llm_function_core.txt (or another appropriate QA
list for multi-node if intended) so it runs in scheduled QA; specifically
include the test class name TestServePrefixAwareScheduling (or its fully
qualified test path if project uses module-style entries) in the
llm_function_core.txt file following the formatting of existing entries.
---
Nitpick comments:
In `@tests/integration/test_lists/test-db/l0_h100.yml`:
- Around line 170-178: Decide whether this new prefix-aware scheduling E2E test
should be run on the scheduled QA list; if it is intended as ongoing release
coverage for core serving behavior, add one representative case from
kv_cache/test_prefix_aware_scheduling.py (for example
TestServePrefixAwareScheduling::test_multi_round_qa_shared_prefix or one
specific parametrized variant like [guaranteed-chunked]) into
tests/integration/test_lists/qa/ so scheduled QA mirrors the pre-merge list; if
it is not intended for scheduled QA, leave the QA list unchanged and document
that the tests are pre-merge only.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0f12cd49-dc28-4b19-91c1-edcf49f850a9
📒 Files selected for processing (7)
cpp/include/tensorrt_llm/batch_manager/templatedTrie.htensorrt_llm/_torch/pyexecutor/scheduler/scheduler.pytests/integration/defs/kv_cache/test_prefix_aware_scheduling.pytests/integration/test_lists/test-db/l0_b200.ymltests/integration/test_lists/test-db/l0_h100.ymltests/unittest/_torch/executor/test_kvcache_aware_router.pytests/unittest/_torch/executor/test_py_scheduler.py
|
PR_Github #45993 [ run ] completed with state |
f7d2e05 to
d6e8f78
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #46442 [ run ] triggered by Bot. Commit: |
|
PR_Github #46442 [ run ] completed with state |
Main failure buckets: FlashInfer / MoE dependency mismatch Affects B300/B200 PyTorch MoE tests and layer-wise benchmark paths. Many H100 failures fail while resolving absolute paths like: DGX_B300-4_GPUs-PyTorch-1: Infra/SSH permission denied, no failed tests. GB200 perf sanity reports a regression, but it is waived and not a blocking unwaived test failure. |
|
/bot run --disable-fail-fast |
|
@tburt-nv |
tburt-nv
left a comment
There was a problem hiding this comment.
The ensure_lmbenchmark fixture looks OK to me, although I'll defer to @NVIDIA/trt-llm-qa regarding the test list changes and the rest of the test code. Unfortunately the new H100 tests didn't run in the latest pipeline due to a timeout, so I've retriggered the pipeline. It should skip the tests that already passed.
|
PR_Github #46530 [ run ] triggered by Bot. Commit: |
|
PR_Github #46530 [ run ] completed with state
|
@tburt-nv They all passes in the new CI. The rest of the failures are unrelated to the code change. I think the PR is safe to merge now. For the report: @schetlur-nv Would you like to review this PR on behalf of @NVIDIA/trt-llm-qa ? Chinese coworkers are in vacation. |
|
/bot skip --comment "CI failures are unrelated. Newly added tests all passed." |
|
PR_Github #46564 [ skip ] triggered by Bot. Commit: |
|
PR_Github #46564 [ skip ] completed with state |
…synthetic-multi-round-qa This is the very first multi-round tests with meaningful prefix reuse. Aim to catch functional bugs with all config combinations of the scheduler. Caught SWA+prefix-aware and python-scheduler bugs during the test developement. Add fixes in the PR. Signed-off-by: Simeng Liu <simengl@nvidia.com>
Signed-off-by: Simeng Liu <simengl@nvidia.com>
d6e8f78 to
ccb8318
Compare
|
/bot skip --comment "Rebase only" |
|
PR_Github #46569 [ skip ] triggered by Bot. Commit: |
|
PR_Github #46569 [ skip ] completed with state |
Signed-off-by: Simeng Liu <simengl@nvidia.com>
…nthetic-multi-round-qa (NVIDIA#13578) Signed-off-by: Simeng Liu <simengl@nvidia.com>
Summary by CodeRabbit
Release Notes
New Features
Bug Fixes
Tests
Description
This is the very first multi-round tests with meaningful prefix reuse.
Aim to catch functional bugs with all config combinations of the
scheduler.
Caught SWA+prefix-aware and python-scheduler bugs during the test
developement.
Add fixes in the PR.
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.