Repository navigation
[Model][PersonaPlex] add per-row RingKV offsets and active rows - #7670
Conversation
Signed-off-by: Mr.Dong <1042542469@qq.com>
…ve-row mask so a row can skip ticks while others advance. Signed-off-by: Mr.Dong <1042542469@qq.com>
Signed-off-by: Mr.Dong <1042542469@qq.com>
Signed-off-by: Mr.Dong <1042542469@qq.com>
MrDongsls
changed the title
Signed-off-by: Mr.Dong <1042542469@qq.com>
Signed-off-by: Mr.Dong <1042542469@qq.com>
Signed-off-by: Mr.Dong <1042542469@qq.com>
MrDongsls
changed the title
|
This PR appears to belong to: docs/design/module/model_integration.md, docs/design/module/ar_runtime.md. Module owners: @tzhouam @gcanlin @fake0fan Routing: @tzhouam via module of the changed files, semantic router, CODEOWNERS; @gcanlin via module of the changed files, semantic router; @fake0fan via module of the changed files @MrDongsls, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer. Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment. |
linyueqian
left a comment
linyueqian
left a comment
There was a problem hiding this comment.
Reviewed at a89c95fb (static read against main plus the CPU and CUDA parity tests as written; nothing executed here). The change is contained to the two PersonaPlex streaming modules and is bit-for-bit the old path when every row is active: with lockstep end_offset the per-row scatter_ hits the same ring slots as the old index_copy_, the position mask is the same formula per row, the RoPE reshape to [B, 1, T, 1] collapses to the old [1, T, 1] broadcast, and _fresh & active is _fresh for an all-true mask, so the first-frame replicate padding and the transposed-conv bias skip are unchanged. active=None maps to the device-resident _all_active, so Stage 0 and Code2Wav keep calling encode_frame / decode_frame exactly as before. Inactive rows are inert where it matters: the gathered K/V slot is written back unchanged, end_offset, prev, partial and _fresh only move where active is true, and reset_slot now masks the row's own write head rather than the shared one. PersonaPlexTemporalStreaming.step has no callers inside vllm_omni yet, so making active required there is safe.
The tests are the right shape: batched rings, transformers and conv carries are compared against singletons that only step on active ticks, and inactive rows are checked against a cloned snapshot; marks match the Simple and Model Executor CUDA steps. One suggestion inline about the mask's device for the future batched callers. Base is a few days old but the merged tree is byte-identical to this head for all three files, and the general lane is green at a89c95fb. Approving.
| def _normalize_temporal_active(active: torch.Tensor, reference: torch.Tensor) -> torch.Tensor: | ||
| if active.shape != reference.shape: | ||
| raise ValueError(f"active must have shape {tuple(reference.shape)}, got {tuple(active.shape)}") | ||
| return active.to(device=reference.device, dtype=torch.bool) |
There was a problem hiding this comment.
[suggestion] Callers are told to pass bool[B], but nothing says which device. A CPU mask lands here in a blocking .to(device=...) copy on every frame (same in _normalize_active in personaplex_mimi.py), which is a host sync per tick and will fail inside CUDA graph capture when the compact batched codec from #7389 gets there. Worth stating in the docstring that the mask must already live on the model device (or accept only that and raise), so the W2 batch builder keeps a persistent device-side mask instead of discovering this under capture.
There was a problem hiding this comment.
Fixed in a follow-up, active is now a persistent device-side bool[B] mask and raises on device or shape mismatches.
…-project#7670) Signed-off-by: Mr.Dong <1042542469@qq.com> Co-authored-by: Mr.Dong <1042542469@qq.com> Signed-off-by: Matthieu Laneuville <matthieu.laneuville@surf.nl>
…-project#7670) Signed-off-by: Mr.Dong <1042542469@qq.com> Co-authored-by: Mr.Dong <1042542469@qq.com>
… out of the per-layer loop Build the per-row RoPE cos/sin tables ([B, 1, T, D/2]) and the RingKV write indexes ([B, T]) / pre-mask position table ([B, capacity]) once per step() in PersonaPlexTemporalStreaming and _MimiStreamingTransformer, and thread them through every layer instead of rebuilding them per layer. The per-row start_offset mask and the active-row write guard stay inside _RingKV.complete(); complete() still builds its own tables when called standalone. Rebased onto the per-row offsets / active-mask API from vllm-project#7670. Signed-off-by: Yifan Tian <Yifan.Tian@colorado.edu>
4 participants
PLEASE FILL IN THE PR DESCRIPTION HERE.
Purpose
W1 standalone enabler described in #7389.
This PR makes PersonaPlex Mimi streaming state row-aware so a compact batched codec can let one row skip a tick without advancing or contaminating its state:
_RingKV.end_offsetand Mimi transformer position progression per-row.active: Tensor[B]mask through Mimi encode/decode so inactive rows preserve their KV, transformer offset, and convolution carry state.Test Plan
vLLM Version:
0.29.0vLLM-Omni Commit:
6c5bd318bb70075dc6f39b141f6178408b58e4e8pytest -q --run-level=core_model -m "core_model and cpu" tests/model_executor/models/personaplex/test_ring_per_row_offsets.py— 4 passedpytest -q --run-level=core_model -m "core_model and cuda" tests/model_executor/models/personaplex/test_ring_per_row_offsets.py- 4 passedpytest -q --run-level=core_model tests/model_executor/models/personaplex --ignore=tests/model_executor/models/personaplex/test_ring_per_row_offsets.py— 41 passedBEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.
(anything written below this line will be removed by GitHub Actions)