[None][fix] Keep MRoPE delta read slots dense across mixed batches - #16944
Merged
longlee0622 merged 2 commits intoJul 29, 2026
Conversation
`_prepare_tp_inputs` only appended to `mrope_delta_read_seq_slots` for generation requests that carried MRoPE metadata, so the gathered delta vector was compacted whenever a request without it shared the batch. The attention kernel indexes `mrope_position_deltas` by *generation batch index* (`decoderMaskedMultiheadAttentionTemplate.h`), so every request scheduled after such a request read another request's delta, and the tail read out of bounds. `cuda_graph_runner` already sizes its static slot buffer as `batch_size * max_beam_width`, i.e. the dense layout was the contract the rest of the stack assumed. Today this is reachable through CUDA graph padding, whose dummy requests have no MRoPE metadata. It becomes reachable for real traffic as soon as an input processor stops synthesizing `mrope_config` for text-only prompts. Resolve requests without MRoPE metadata to the reserved zero slot instead of dropping them: their delta is zero by construction, and the reserved slot permanently reads back zero because the cache is zero-initialized and the write path only targets real `py_seq_slot`s. When no generation request in the batch has MRoPE metadata the list is dropped entirely, so the all-zero vector does not block the steady-state generation fast path. Also stop indexing `mrope_config` and `mrope_position_ids` directly on the context and generation paths, so a request may carry multimodal content without MRoPE metadata. Verified on Qwen2.5-VL-7B: a batch of image requests interleaved with text-only requests whose input processor emits no `mrope_config` fails with `cudaErrorIllegalAddress` before this change and reproduces the multimodal-only baseline outputs after it. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Without vision spans the M-RoPE coordinates degenerate to the scalar token positions on all three axes and the position delta is zero, which is what the model engine already falls back to for a request carrying no `mrope_config`. Synthesizing them anyway costs an O(seq_len) (3, 1, N) int32 tensor per request plus its H2D copy, and in disaggregated serving the prefill worker re-registers that tensor as a CUDA IPC handle that nothing downstream ever consumes. Return no extra processed inputs from the text-only fast path instead, so text traffic on an M-RoPE model allocates nothing and no longer builds a `MultimodalParams` at all. Also refresh the two comments that documented the old "every Qwen*-VL request carries mrope_config" invariant. Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
yechank-nvidia
requested review from
2ez4bz,
JunyiXu-nv,
asfiyab-nvidia,
bo-nv and
chienchunhung
July 28, 2026 09:19
Contributor
WalkthroughText-only Qwen2VL-family prompts now omit multimodal and MRoPE metadata. MRoPE input preparation handles mixed batches with dense dummy slots and removes unused delta slots for all-text generation batches, with expanded processor and executor tests. ChangesMRoPE text-only and mixed-batch handling
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested labels: 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}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_pytorch_model_engine.py (1)
610-649: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDeduplicate MRoPE engine setup between the old and new tests.
_setup_mrope_engine(Lines 699-738) is a near-verbatim extraction of the inline setup already present intest_prepare_tp_inputs_with_partial_mrope_segments(Lines 610-649): samellm_args/DummyModelEngineconstruction,Mapping,KvCacheConfig/KVCacheManager,AttentionMetadata, and CUDA buffer allocations. Since the helper now exists, refactor the older test to call it too, rather than keeping two copies of the same setup in the file.♻️ Proposed refactor
def test_prepare_tp_inputs_with_partial_mrope_segments(self) -> None: """Test generation-only MRoPE assembly with a real multimodal span and a dummy padded request.""" - llm_args = TorchLlmArgs(model="dummy") - model_engine = DummyModelEngine(llm_args, dtype=torch.half) - model_engine.model.model_config.pretrained_config.rope_scaling = { - "type": "mrope" - } - - mapping = Mapping(world_size=1, tp_size=1, rank=0) - kv_cache_config = KvCacheConfig(max_tokens=32) - kv_cache_manager = KVCacheManager( - kv_cache_config, - tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=1, - num_kv_heads=16, - head_dim=16, - tokens_per_block=1, - max_seq_len=32, - max_batch_size=4, - mapping=mapping, - dtype=tensorrt_llm.bindings.DataType.HALF, - ) - attn_metadata = AttentionMetadata(max_num_requests=4, - max_num_tokens=32, - kv_cache_manager=kv_cache_manager) - attn_metadata.is_cuda_graph = False - - model_engine.max_num_tokens = 32 - model_engine.input_ids_cuda = torch.zeros(32, - dtype=torch.int32, - device='cuda') - model_engine.position_ids_cuda = torch.zeros(32, - dtype=torch.int32, - device='cuda') - model_engine.mrope_position_ids_cuda = torch.zeros((3, 1, 32), - dtype=torch.int32, - device='cuda') - model_engine.previous_batch_indices_cuda = torch.zeros( - 32, dtype=torch.int32, device='cuda') + model_engine, kv_cache_manager, attn_metadata = self._setup_mrope_engine( + )Also applies to: 699-738
🤖 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/_torch/executor/test_pytorch_model_engine.py` around lines 610 - 649, Refactor test_prepare_tp_inputs_with_partial_mrope_segments to use the existing _setup_mrope_engine helper instead of duplicating llm_args, DummyModelEngine, cache-manager, attention-metadata, and CUDA-buffer setup. Preserve the test’s MRoPE configuration and any helper-returned values needed by the remaining assertions.
🤖 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.
Nitpick comments:
In `@tests/unittest/_torch/executor/test_pytorch_model_engine.py`:
- Around line 610-649: Refactor
test_prepare_tp_inputs_with_partial_mrope_segments to use the existing
_setup_mrope_engine helper instead of duplicating llm_args, DummyModelEngine,
cache-manager, attention-metadata, and CUDA-buffer setup. Preserve the test’s
MRoPE configuration and any helper-returned values needed by the remaining
assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2c1c4af6-cbee-48ae-a677-550684047d9a
📒 Files selected for processing (5)
tensorrt_llm/_torch/models/modeling_qwen2vl.pytensorrt_llm/_torch/models/modeling_qwen3vl.pytensorrt_llm/_torch/pyexecutor/model_engine.pytests/unittest/_torch/executor/test_pytorch_model_engine.pytests/unittest/_torch/multimodal/test_qwen2vl_text_only_prompt.py
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #62169 [ run ] triggered by Bot. Commit: |
Shixiaowei02
approved these changes
Jul 28, 2026
Collaborator
|
PR_Github #62169 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #62204 [ run ] triggered by Bot. Commit: |
2ez4bz
approved these changes
Jul 28, 2026
Collaborator
|
PR_Github #62204 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62351 [ run ] triggered by Bot. Commit: |
longlee0622
enabled auto-merge (squash)
July 29, 2026 02:48
JunyiXu-nv
approved these changes
Jul 29, 2026
JunyiXu-nv
left a comment
JunyiXu-nv
left a comment
Collaborator
There was a problem hiding this comment.
LGTM from runtime side!
Collaborator
|
PR_Github #62351 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62433 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62433 [ run ] completed with state |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
_prepare_tp_inputsappended tomrope_delta_read_seq_slotsonly forgeneration requests carrying MRoPE metadata, compacting the delta vector
whenever a request without it shared the batch. The attention kernel indexes
mrope_position_deltasby generation batch index(
decoderMaskedMultiheadAttentionTemplate.h), so every later request readanother request's delta and the tail read out of bounds — silent corruption,
no shape check anywhere.
cuda_graph_runneralready sizes its static slotbuffer as
batch_size * max_beam_width, so dense was the assumed contract.Reachable on
maintoday via CUDA graph padding (dummy requests have no MRoPEmetadata); dummy outputs are discarded so it went unnoticed. It hits real
traffic as soon as an input processor stops synthesizing
mrope_configfortext-only prompts.
Commit 1: requests without MRoPE metadata resolve to the reserved zero slot
instead of being dropped — their delta is zero by construction, and that slot
always reads back zero. If no generation request has MRoPE metadata the list is
dropped entirely, keeping the steady-state generation fast path reachable.
Commit 2: with the runtime no longer requiring it, the text-only fast path stops
synthesizing MRoPE metadata, saving an O(seq_len)
(3, 1, N)tensor plus H2Dper request — and in disagg a CUDA IPC handle nothing consumes.
Text-only observation raised in #16874; this lands the runtime fix it depends on.
Test Coverage
test_pytorch_model_engine.py: 2 new tests (dense slots in a mixed batch;all-text-only batch drops the deltas) +
test_prepare_tp_inputs_with_partial_mrope_segmentsupdated — the padded dummy now yields
[0, 32]instead of[0].test_qwen2vl_text_only_prompt.py(new): 9 cases across the three Qwen VLprocessors,
multi_modal_dataabsent /None/ empty.Local Qwen2.5-VL-7B: interleaving text-only with image requests fails with
cudaErrorIllegalAddresswithout the runtime fix, matches the multimodal-onlybaseline with it.
Dev Engineer Review
QA Engineer Review
test_prepare_tp_inputs_mixed_text_only_keeps_mrope_deltas_densetest_prepare_tp_inputs_all_text_only_drops_mrope_deltastest_text_only_prompt_emits_no_multimodal_datatest_text_only_prompt_without_multi_modal_data_keytests/integration/test_lists/test-db or QA lists.