[None][fix] Release stale-range holder after _commit_block in KVCache… - #14478
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #50027 [ run ] triggered by Bot. Commit: |
|
PR_Github #50027 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #50071 [ run ] triggered by Bot. Commit: |
|
PR_Github #50071 [ run ] completed with state |
|
Could we add a focused regression UT for this state transition? I locally reproduced the pre-fix failure on a B200 node with a minimal KVCacheManagerV2 sequence:
The observed state was:
So my read is that existing SWA tests cover the normal commit-before-stale path and the scratch-reuse path, but not this history-before-commit path. A focused test in |
|
For unit test, I think we can cover it by updating existing tests instead of creating a new one. We can add a parameter delay_commit to TestNoBatching.run_naive defaulting to False, then if True, only commit all tokens at end. And add that arg to TestNoBatching.test_naive(). |
Address review comment on PR NVIDIA#14478: use the life_cycles.attention_life_cycles() iterator (which already filters out the SSM lifecycle) instead of items() + a manual ssm_life_cycle_id guard, drop the redundant 'is not None' check (assigning None to None is a no-op), and swap the loop nesting so the per-lifecycle stale-range computation runs once per lifecycle rather than once per beam x lifecycle. Behavior is unchanged: for every committed seq_block whose ordinal sits in the stale range of an attention lifecycle, the per-beam holder for that lifecycle is released, restoring the 'committed AND stale -> holder is None' invariant that _check_sanity expects in suspend(). Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
…ManagerV2
When _commit_block finalizes a block that is already inside stale_range[lc]
for some finite-window lifecycle (e.g. DSv4 DSA / SWA), the freshly assigned
holder (_PageHolder or _SharedPageLock from line 1268 / 1312-1325) was left
in place. _check_sanity then asserted on the precondition of suspend():
/workspaces/tensorrt_llm/.../sampler.../suspend()
-> _kv_cache.py:931 assert self._check_sanity()
-> _kv_cache.py:1531 assert holder is None # for committed stale block
This trips deterministically on DSv4-Pro disagg GEN runs under MAX_UTILIZATION
scheduling: when the V2 scheduler tries to suspend an in-flight request to
make room, the precondition check aborts the worker. The previously-known
release paths (_on_stop_committing, _unlock_stale_blocks) never revisit a
committed-and-already-stale ordinal, so a 'commit-into-stale-range' case
slipped through.
Fix: at the tail of _commit_block, for the committed block, iterate each
non-SSM lifecycle and clear beam_block[lc] when the ordinal lies inside
that lifecycle's stale_range. The committed page itself remains alive in
the radix tree (committed pages are shareable across caches); only this
_KVCache's redundant strong reference is dropped, restoring the
'committed AND stale -> holder is None' invariant.
Validated end-to-end on DSv4-Pro disagg sweep on Lyris GB300 (c=4/8/16/32/64,
DEP8/TEP8): zero _check_sanity AssertionError across ~160k commit-into-stale
events under real traffic.
Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
…ng.test_naive Add a delay_commit parameter to TestNoBatching.run_request / TestNoBatching.run_naive that, when True, skips the per-step kv_cache.commit() inside the decode loop and only commits the full generated range at the very end of the request. This exercises the "commit a large pending range into a stale lifecycle window" path that the preceding fix targets (commit-into-stale-range where _commit_block needs to release a holder it just placed). Extend TestNoBatching.test_naive to parameterize over delay_commit alongside the existing use_external_page_index_buf x use_block_quant axes; this goes from 4 to 8 parametrized cases with no new test function. All other call sites of run_request / run_naive default to delay_commit=False and keep their existing behavior unchanged. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
Address review comment on PR NVIDIA#14478: use the life_cycles.attention_life_cycles() iterator (which already filters out the SSM lifecycle) instead of items() + a manual ssm_life_cycle_id guard, drop the redundant 'is not None' check (assigning None to None is a no-op), and swap the loop nesting so the per-lifecycle stale-range computation runs once per lifecycle rather than once per beam x lifecycle. Behavior is unchanged: for every committed seq_block whose ordinal sits in the stale range of an attention lifecycle, the per-beam holder for that lifecycle is released, restoring the 'committed AND stale -> holder is None' invariant that _check_sanity expects in suspend(). Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
390c5d5 to
e0e3ce0
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #50406 [ run ] triggered by Bot. Commit: |
|
PR_Github #50406 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #50583 [ run ] triggered by Bot. Commit: |
|
PR_Github #50583 [ run ] completed with state |
NVIDIA#14478) Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
NVIDIA#14478) Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
NVIDIA#14478) Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com>
5 participants
…ManagerV2
When _commit_block finalizes a block that is already inside stale_range[lc] for some finite-window lifecycle (e.g. DSv4 DSA / SWA), the freshly assigned holder (_PageHolder or _SharedPageLock from line 1268 / 1312-1325) was left in place. _check_sanity then asserted on the precondition of suspend():
/workspaces/tensorrt_llm/.../sampler.../suspend()
-> _kv_cache.py:931 assert self._check_sanity()
-> _kv_cache.py:1531 assert holder is None # for committed stale block
This trips deterministically on DSv4-Pro disagg GEN runs under MAX_UTILIZATION scheduling: when the V2 scheduler tries to suspend an in-flight request to make room, the precondition check aborts the worker. The previously-known release paths (_on_stop_committing, _unlock_stale_blocks) never revisit a committed-and-already-stale ordinal, so a 'commit-into-stale-range' case slipped through.
Fix: at the tail of _commit_block, for the committed block, iterate each non-SSM lifecycle and clear beam_block[lc] when the ordinal lies inside that lifecycle's stale_range. The committed page itself remains alive in the radix tree (committed pages are shareable across caches); only this _KVCache's redundant strong reference is dropped, restoring the 'committed AND stale -> holder is None' invariant.
Validated end-to-end on DSv4-Pro disagg sweep on Lyris GB300 (c=4/8/16/32/64, DEP8/TEP8): zero _check_sanity AssertionError across ~160k commit-into-stale events under real traffic.
@coderabbitai summary
Description
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)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.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.