[None][fix] disagg: derive KV transfer layer offset from physical slot order - #16429
Conversation
📝 WalkthroughWalkthrough
ChangesKV map alignment
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![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 (1)
tests/unittest/disaggregated/test_peer.py (1)
407-456: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCoverage needs follow-up for the other new branches.
Current coverage is insufficient: it only tests INDEXED ordering and local-side non-contiguity. Please add:
test_peer_registrar_get_kv_map_uses_physical_offset_order_flattest_peer_registrar_get_kv_map_rejects_peer_non_contiguous_overlapor a reordered peer run caseRun
pytest tests/unittest/disaggregated/test_peer.py -k get_kv_map.As per path instructions, provide concrete test names and state whether TensorRT-LLM test coverage is sufficient.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/disaggregated/test_peer.py` around lines 407 - 456, Add tests named test_peer_registrar_get_kv_map_uses_physical_offset_order_flat and test_peer_registrar_get_kv_map_rejects_peer_non_contiguous_overlap, covering flat physical ordering and a peer-side non-contiguous or reordered overlap rejection; keep the existing INDEXED and local-side coverage intact. Run pytest tests/unittest/disaggregated/test_peer.py -k get_kv_map and state whether TensorRT-LLM test coverage is sufficient.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/disaggregated/test_peer.py`:
- Line 407: Add return annotations of -> None to both test functions,
test_peer_registrar_get_kv_map_uses_physical_offset_order and the additional
test at the referenced location, without changing their behavior.
---
Nitpick comments:
In `@tests/unittest/disaggregated/test_peer.py`:
- Around line 407-456: Add tests named
test_peer_registrar_get_kv_map_uses_physical_offset_order_flat and
test_peer_registrar_get_kv_map_rejects_peer_non_contiguous_overlap, covering
flat physical ordering and a peer-side non-contiguous or reordered overlap
rejection; keep the existing INDEXED and local-side coverage intact. Run pytest
tests/unittest/disaggregated/test_peer.py -k get_kv_map and state whether
TensorRT-LLM test coverage is sufficient.
🪄 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: fb2df76a-e75e-472f-97f6-70571f7452cf
📒 Files selected for processing (2)
tensorrt_llm/_torch/disaggregation/native/peer.pytests/unittest/disaggregated/test_peer.py
|
/bot run --disable-fail-fast |
|
PR_Github #59456 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #59487 [ run ] triggered by Bot. Commit: |
|
PR_Github #59456 [ run ] completed with state |
|
PR_Github #59487 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59638 [ run ] triggered by Bot. Commit: |
|
PR_Github #59638 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59776 [ run ] triggered by Bot. Commit: |
|
PR_Github #59776 [ run ] completed with state
|
…t order PeerRegistrar.get_kv_map() sorted the pool's global layer ids before computing each layer's slot offset, relying on the convention that managers assign global_layer_id monotonically with the layer's byte offset. That discards the physical-offset order recovered by get_pool_view_global_layer_ids() and, for a non-monotonic layout, maps a layer to the wrong slot and transfers the wrong KV bytes. Use the physical slot order directly (buffer_entries offsets for INDEXED pools, local_layers packing order for FLAT pools) so a layer's index is its slot offset, and anchor the transfer on the physically-first overlapping layer. Replace the implicit monotonicity invariant with an explicit check that the shared layers form an aligned contiguous slot run on both peers, raising instead of silently corrupting the transfer. Add unit tests covering a non-monotonic layout and the non-contiguous-overlap rejection. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Address review nit: branch explicitly on MapperKind.INDEXED instead of relying on else, and raise ValueError for unexpected mapper kinds. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
cb50d56 to
d7c7030
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59968 [ run ] triggered by Bot. Commit: |
|
PR_Github #59968 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59999 [ run ] triggered by Bot. Commit: |
|
PR_Github #59999 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60027 [ run ] triggered by Bot. Commit: |
|
PR_Github #60027 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60086 [ run ] triggered by Bot. Commit: |
|
PR_Github #60086 [ run ] completed with state |
3 participants
Description
Follow-up to a review comment on #16218 (discussion).
PeerRegistrar.get_kv_map()sorted the pool's global layer ids byglobal_layer_idbefore computing each layer's slot offset, relying on the convention that KV-cache managers assignglobal_layer_idmonotonically with the layer's byte offset in the slot. That discards the physical-offset order thatget_pool_view_global_layer_ids()already recovers from the V2 pool-viewbuffer_entries. For a non-monotonic layout (e.g. physical order[5, 3]), a layer is mapped to the wrong slot and the transfer copies the wrong KV bytes.This was not a regression — the
sorted()and the monotonicity assumption pre-existed onmain, and under that assumption sorted order equals physical order — so it is fixed here in a dedicated PR rather than in the test-prep change #16218.What changed
get_kv_map(), droppingsorted()on both branches:get_pool_view_global_layer_ids()— ordered bybuffer_entriesoffsets (the V2 pool-view layout).get_global_layer_ids()—local_layerspacking order, which is the physical order for FLAT pools (previously also wronglysorted())..index()yields the true slot offset.ValueErrorinstead of silently transferring the wrong bytes.Tests
test_peer_registrar_get_kv_map_uses_physical_offset_order: non-monotonic layout[5, 3]; asserts the mapper uses the physical slot offset (1), which the old sorted path got wrong.test_peer_registrar_get_kv_map_rejects_non_contiguous_overlap: verifies the new contiguity guard raises.Test Coverage
tests/unittest/disaggregated/test_peer.pyPR Checklist
Summary by CodeRabbit
Bug Fixes
Tests