[TRTLLM-9409][feat] Pass MRoPE tensors for EPD disagg - #9758
Conversation
671c8c1 to
365eb1c
Compare
365eb1c to
a553c4b
Compare
jaedeok-nvidia
left a comment
jaedeok-nvidia
left a comment
There was a problem hiding this comment.
Many thanks @2ez4bz for adding mrope support for multimodal EPD disagg serving. I've left several questions. Would you please take a look? In overall it looks good to me, although the detailed logic should be confirmed by Yechan and Chang.
a553c4b to
82a0a48
Compare
82a0a48 to
0c02092
Compare
yechank-nvidia
left a comment
yechank-nvidia
left a comment
There was a problem hiding this comment.
LGTM. Thx for the work!
a956430 to
c770297
Compare
2ez4bz
changed the title
📝 WalkthroughWalkthroughThis change introduces support for disaggregated multimodal handling with MRope position data propagation. New Qwen2.5 VL model support is added, multimodal encoders can be disabled during disaggregation, and MRope metadata is threaded through the executor pipeline with conditional device transfers. Test coverage is refactored to use direct model directory inputs instead of model-key-based lookups. Changes
Sequence Diagram(s)sequenceDiagram
participant Executor
participant ModelEngine
participant Sampler
participant PyResult
Executor->>ModelEngine: _prepare_tp_inputs(multimodal_params with mrope)
alt CPU-resident mrope_position_deltas
ModelEngine->>ModelEngine: to_device(multimodal_data)
else GPU-resident mrope_position_deltas
ModelEngine->>ModelEngine: skip transfer
end
ModelEngine->>ModelEngine: _forward_step_mm_encoder_only()
ModelEngine->>Sampler: return mm_embeddings + mrope_position_ids/deltas
Sampler->>Sampler: sample_async: wrap in MultimodalResult
Sampler->>Sampler: pack mrope metadata into extra_data
Sampler->>PyResult: update_requests: extract extra_data
PyResult->>PyResult: set_mrope_position(ids, deltas)
PyResult->>Executor: return LlmResult with mrope handles
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Areas requiring extra attention:
Pre-merge checks and finishing touches❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 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: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (7)
tensorrt_llm/_torch/pyexecutor/sampler.py (1)
258-271: Avoid mutatingmodel_outputsviapop(); buildextra_datawithout side effects.
model_outputs.pop("mm_embeddings")can break callers if they reusemodel_outputsafter sampling.@@ - data = MultimodalResult( - mm_embeddings=model_outputs.pop("mm_embeddings"), - extra_data={**model_outputs}, - ) + mm_embeddings = model_outputs["mm_embeddings"] + extra_data = {k: v for k, v in model_outputs.items() if k != "mm_embeddings"} + data = MultimodalResult( + mm_embeddings=mm_embeddings, + extra_data=(extra_data or None), + )tensorrt_llm/executor/result.py (1)
438-456: Guard optional mRoPE fields + fix tuple assignment tomultimodal_embedding_handles.
response_result.mrope_position_*_handleshould be accessed viagetattr(like mm_embedding_handle) to avoid AttributeError on models that don’t return it.- The trailing comma makes
multimodal_embedding_handlesa tuple.@@ if hasattr(response_result, 'mm_embedding_handle' ) and response_result.mm_embedding_handle is not None: self._mm_embedding_handle = response_result.mm_embedding_handle - mrope_position_ids_handle = response_result.mrope_position_ids_handle - mrope_position_deltas_handle = response_result.mrope_position_deltas_handle + mrope_position_ids_handle = getattr(response_result, "mrope_position_ids_handle", None) + mrope_position_deltas_handle = getattr(response_result, "mrope_position_deltas_handle", None) if self.disaggregated_params is not None: - self.disaggregated_params.multimodal_embedding_handles = [ - response_result.mm_embedding_handle - ], + self.disaggregated_params.multimodal_embedding_handles = [ + response_result.mm_embedding_handle + ] self.disaggregated_params.multimodal_hashes = self._multimodal_hashes else: self.disaggregated_params = DisaggregatedParams( multimodal_embedding_handles=[ response_result.mm_embedding_handle ], multimodal_hashes=self._multimodal_hashes) - self.disaggregated_params.mrope_position_ids_handle = mrope_position_ids_handle - self.disaggregated_params.mrope_position_deltas_handle = mrope_position_deltas_handle + if mrope_position_ids_handle is not None or mrope_position_deltas_handle is not None: + self.disaggregated_params.mrope_position_ids_handle = mrope_position_ids_handle + self.disaggregated_params.mrope_position_deltas_handle = mrope_position_deltas_handletensorrt_llm/_torch/pyexecutor/model_engine.py (3)
1-1: Add missing NVIDIA copyright header (current year) at file top.Per coding guidelines for
**/*.{cpp,h,cu,py}, this file should have the NVIDIA header (with 2025) before imports.
1938-1946: Potential crash:torch.cat(mrope_position_ids, ...)can see mixed CPU/CUDA tensors.With mixed scheduling (some requests already in generation with recovered CUDA MRoPE tensors, others new context with CPU MRoPE tensors),
torch.cat(...)will throw a device mismatch.One low-impact fix is to normalize all
mrope_position_idsto a single device before concatenation (prefer CUDA if any are CUDA):- mrope_position_ids = torch.cat(mrope_position_ids, dim=-1) - if mrope_position_ids.device.type == "cpu": - mrope_position_ids = mrope_position_ids.pin_memory() + # Normalize to a single device to avoid cat() device-mismatch in mixed batches. + if any(t.device.type == "cuda" for t in mrope_position_ids): + target_device = self.mrope_position_ids_cuda.device + normalized = [] + for t in mrope_position_ids: + if t.device != target_device: + if t.device.type == "cpu": + t = t.pin_memory().to(target_device, non_blocking=True) + else: + t = t.to(target_device) + normalized.append(t) + mrope_position_ids = torch.cat(normalized, dim=-1) + else: + mrope_position_ids = torch.cat(mrope_position_ids, dim=-1).pin_memory()
2805-2814: Keep mm-encoder-only output schema consistent (always includelogits).You now return
{'mm_embeddings': ..., 'logits': None, ...}but the “no multimodal data” early-return still returns only{'mm_embeddings': []}. Suggest:- return {'mm_embeddings': []} + return {'mm_embeddings': [], 'logits': None}Also applies to: 2840-2860
tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py (1)
1-1: Add missing NVIDIA copyright header (current year) at file top.Per coding guidelines for
**/*.{cpp,h,cu,py}, tests should also carry the NVIDIA header (with 2025).tensorrt_llm/_torch/models/modeling_qwen2vl.py (1)
1-1: Add missing NVIDIA copyright header (current year) at file top.Per coding guidelines for
**/*.{cpp,h,cu,py}, this file should have the NVIDIA header (with 2025) before imports.
🧹 Nitpick comments (4)
tensorrt_llm/disaggregated_params.py (1)
43-44: Consider adding validation in__post_init__for consistency.The code in
llm.py(lines 447, 467-470) enforces that bothmrope_position_ids_handleandmrope_position_deltas_handlemust be provided together or both beNone. Consider adding this validation here in__post_init__to fail fast with a clear error message, similar to the existing validation formultimodal_embedding_handlesandmultimodal_hashes.def __post_init__(self): if self.request_type is not None: self.request_type = self.request_type.lower() if self.request_type not in [ "context_only", "generation_only", "context_and_generation", ]: raise ValueError( f"Unknown request type: {self.request_type}. Must be context_only, generation_only or " "context_and_generation" ) + # Validate mrope handles are both present or both None + if (self.mrope_position_ids_handle is None) != (self.mrope_position_deltas_handle is None): + raise ValueError( + "mrope_position_ids_handle and mrope_position_deltas_handle must both be provided or both be None" + ) if self.multimodal_embedding_handles is not None:tensorrt_llm/_torch/pyexecutor/sampler.py (1)
199-207: PublicMultimodalResult.extra_datashould be documented via docstring/field doc, not a line comment.Per coding guidelines (“prefer docstrings over comments” for public interfaces), consider turning the comment into a class docstring or a field doc block.
tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py (1)
126-130: Avoid mutatinginputsin-place for the decode step (test robustness).Consider creating a shallow copy for the decode call so future assertions/extensions don’t accidentally depend on the mutated state:
- inputs[0]['multi_modal_data'] = None - inputs[0]['prompt_token_ids'] = outputs[0].prompt_token_ids + decode_inputs = [dict(inputs[0])] + decode_inputs[0]['multi_modal_data'] = None + decode_inputs[0]['prompt_token_ids'] = outputs[0].prompt_token_ids ... - outputs = llm_decode.generate( - inputs, + outputs = llm_decode.generate( + decode_inputs, sampling_params=sampling_params, disaggregated_params=pd_disaggregated_params)tensorrt_llm/_torch/models/modeling_qwen2vl.py (1)
1051-1126: Hardenget_prompt_token_ids(avoidassertfor user/input validation; address TRY003).Assertions can be optimized out; for input validation, explicit exceptions are safer. Also ruff TRY003 is currently flagging the raised exceptions.
- hidden_size = mm_handles[0]['tensor_size'][1] - assert hidden_size == self.config.text_config.hidden_size, "Multimodal embedding hidden size must match model hidden size" + hidden_size = mm_handles[0]["tensor_size"][1] + if hidden_size != self.config.text_config.hidden_size: + raise ValueError("Multimodal embedding hidden size must match model hidden size") ... - assert write_pos == final_length, f"Write position mismatch: {write_pos} != {final_length}" - assert mm_token_length[-1] + mm_token_offsets[ - -1] <= final_length, f"mm_token_length[-1] + mm_token_offsets[-1] ({mm_token_length[-1] + mm_token_offsets[-1]}) should be less than or equal to final_length ({final_length})" + if write_pos != final_length: + raise RuntimeError(f"Write position mismatch: {write_pos} != {final_length}") + if mm_token_length[-1] + mm_token_offsets[-1] > final_length: + raise RuntimeError("Expanded multimodal placeholders exceed final token length")
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
tensorrt_llm/_torch/models/modeling_llava_next.py(1 hunks)tensorrt_llm/_torch/models/modeling_qwen2vl.py(4 hunks)tensorrt_llm/_torch/pyexecutor/llm_request.py(4 hunks)tensorrt_llm/_torch/pyexecutor/model_engine.py(3 hunks)tensorrt_llm/_torch/pyexecutor/sampler.py(5 hunks)tensorrt_llm/disaggregated_params.py(1 hunks)tensorrt_llm/executor/result.py(5 hunks)tensorrt_llm/llmapi/llm.py(3 hunks)tensorrt_llm/llmapi/mm_encoder.py(1 hunks)tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py(6 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+
Indent Python code with 4 spaces; do not use tabs
Always maintain the namespace when importing in Python, even if only one class or function from a module is used (e.g., usefrom package.subpackage import fooand thenfoo.SomeClass()instead offrom package.subpackage.foo import SomeClass)
Python filenames should use snake_case (e.g.,some_file.py)
Python class names should use PascalCase (e.g.,class SomeClass)
Python function and method names should use snake_case (e.g.,def my_awesome_function():)
Python local variable names should use snake_case, with prefixkfor variable names that start with a number (e.g.,k_99th_percentile = ...)
Python global variables should use upper snake_case with prefixG(e.g.,G_MY_GLOBAL = ...)
Python constants should use upper snake_case (e.g.,MY_CONSTANT = ...)
Avoid shadowing variables declared in an outer scope in Python
Initialize all externally visible members of a Python class in the constructor
For Python interfaces that may be used outside a file, prefer docstrings over comments
Python comments should be reserved for code within a function, or interfaces that are local to a file
Use Google style docstrings for Python classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with type and description (e.g.,self.x = 5followed by"""<type>: Description of 'x'""")
Avoid using reflection in Python when functionality can be easily achieved without reflection
When using try-except blocks in Python, limit the except clause to the smallest set of specific errors possible instead of catching all exceptions
When using try-except blocks in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible and use the else block to implement the logic
Files:
tensorrt_llm/disaggregated_params.pytensorrt_llm/llmapi/mm_encoder.pytensorrt_llm/_torch/models/modeling_llava_next.pytensorrt_llm/_torch/pyexecutor/sampler.pytensorrt_llm/_torch/pyexecutor/llm_request.pytests/unittest/_torch/multimodal/test_mm_encoder_standalone.pytensorrt_llm/llmapi/llm.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/models/modeling_qwen2vl.pytensorrt_llm/executor/result.py
**/*.{cpp,h,cu,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code files should contain an NVIDIA copyright header that includes the current year at the top
Files:
tensorrt_llm/disaggregated_params.pytensorrt_llm/llmapi/mm_encoder.pytensorrt_llm/_torch/models/modeling_llava_next.pytensorrt_llm/_torch/pyexecutor/sampler.pytensorrt_llm/_torch/pyexecutor/llm_request.pytests/unittest/_torch/multimodal/test_mm_encoder_standalone.pytensorrt_llm/llmapi/llm.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/models/modeling_qwen2vl.pytensorrt_llm/executor/result.py
🧠 Learnings (14)
📚 Learning: 2025-08-14T15:38:01.771Z
Learnt from: MatthiasKohl
Repo: NVIDIA/TensorRT-LLM PR: 6904
File: cpp/tensorrt_llm/pybind/thop/bindings.cpp:55-57
Timestamp: 2025-08-14T15:38:01.771Z
Learning: In TensorRT-LLM Python bindings, tensor parameter collections like mla_tensor_params and spec_decoding_tensor_params are kept as required parameters without defaults to maintain API consistency, even when it might affect backward compatibility.
Applied to files:
tensorrt_llm/disaggregated_params.pytensorrt_llm/_torch/pyexecutor/sampler.pytensorrt_llm/llmapi/llm.py
📚 Learning: 2025-07-22T09:22:14.726Z
Learnt from: yechank-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.726Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
Applied to files:
tensorrt_llm/llmapi/mm_encoder.pytensorrt_llm/llmapi/llm.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/executor/result.py
📚 Learning: 2025-08-14T06:36:40.701Z
Learnt from: timlee0212
Repo: NVIDIA/TensorRT-LLM PR: 6886
File: tensorrt_llm/_torch/models/modeling_deepseekv3.py:0-0
Timestamp: 2025-08-14T06:36:40.701Z
Learning: In DeepSeek V3 model (tensorrt_llm/_torch/models/modeling_deepseekv3.py), the disagreement between AllReduce.__init__ guard and _compute_mlp_tp_size logic for MNNVL usage is expected by design. The AllReduce component and MLP TP-size computation intentionally use different criteria for MNNVL availability decisions.
Applied to files:
tensorrt_llm/_torch/models/modeling_llava_next.pytensorrt_llm/llmapi/llm.pytensorrt_llm/_torch/models/modeling_qwen2vl.py
📚 Learning: 2025-07-28T17:06:08.621Z
Learnt from: moraxu
Repo: NVIDIA/TensorRT-LLM PR: 6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.
Applied to files:
tests/unittest/_torch/multimodal/test_mm_encoder_standalone.pytensorrt_llm/llmapi/llm.py
📚 Learning: 2025-08-11T20:09:24.389Z
Learnt from: achartier
Repo: NVIDIA/TensorRT-LLM PR: 6763
File: tests/integration/defs/triton_server/conftest.py:16-22
Timestamp: 2025-08-11T20:09:24.389Z
Learning: In the TensorRT-LLM test infrastructure, the team prefers simple, direct solutions (like hard-coding directory traversal counts) over more complex but robust approaches when dealing with stable directory structures. They accept the maintenance cost of updating tests if the layout changes.
Applied to files:
tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
📚 Learning: 2025-09-09T09:40:45.658Z
Learnt from: fredricz-20070104
Repo: NVIDIA/TensorRT-LLM PR: 7645
File: tests/integration/test_lists/qa/llm_function_core.txt:648-648
Timestamp: 2025-09-09T09:40:45.658Z
Learning: In TensorRT-LLM test lists, it's common and intentional for the same test to appear in multiple test list files when they serve different purposes (e.g., llm_function_core.txt for comprehensive core functionality testing and llm_function_core_sanity.txt for quick sanity checks). This duplication allows tests to be run in different testing contexts.
Applied to files:
tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
📚 Learning: 2025-09-03T13:16:38.028Z
Learnt from: nvpohanh
Repo: NVIDIA/TensorRT-LLM PR: 7478
File: tests/unittest/_torch/modeling/test_modeling_llama_min_latency.py:286-308
Timestamp: 2025-09-03T13:16:38.028Z
Learning: In test files, temporary monkey-patches for upstream bugs can be kept simple when they are explicitly intended to be removed soon, rather than investing effort in making them more robust.
Applied to files:
tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
📚 Learning: 2025-08-20T06:56:02.889Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:577-579
Timestamp: 2025-08-20T06:56:02.889Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, maxSequenceLength is now enforced as a non-optional argument in the BlockManager constructor, so concerns about std::nullopt defaulting to 0 are not applicable. When windowSize > maxSequenceLength, a warning should be added instead of handling optional parameter cases.
Applied to files:
tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
📚 Learning: 2025-08-13T08:21:13.006Z
Learnt from: dbari
Repo: NVIDIA/TensorRT-LLM PR: 6714
File: tests/integration/defs/triton_server/build_model.sh:724-726
Timestamp: 2025-08-13T08:21:13.006Z
Learning: Mistral Small 3.1 multimodal (pixtral model type) supports any batch size for the multimodal encoder, not just batch size 1. The max_batch_size parameter can be set to values like 2 without conflicts with the runtime batching logic.
Applied to files:
tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py
📚 Learning: 2025-08-21T02:39:12.009Z
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 7104
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1475-1480
Timestamp: 2025-08-21T02:39:12.009Z
Learning: The min latency mode functionality in TensorRT-LLM MOE kernels (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu) is deprecated and no longer being maintained/updated, as confirmed by djns99. Bug reports and optimization suggestions for the computeStridesTmaWarpSpecializedLowLatencyKernel and related min latency code paths should be deprioritized.
Applied to files:
tensorrt_llm/llmapi/llm.py
📚 Learning: 2025-08-14T15:43:23.107Z
Learnt from: MatthiasKohl
Repo: NVIDIA/TensorRT-LLM PR: 6904
File: tensorrt_llm/_torch/attention_backend/trtllm.py:259-262
Timestamp: 2025-08-14T15:43:23.107Z
Learning: In TensorRT-LLM's attention backend, tensor parameters in the plan() method are assigned directly without validation (dtype, device, contiguity checks). This maintains consistency across all tensor inputs and follows the pattern of trusting callers to provide correctly formatted tensors.
Applied to files:
tensorrt_llm/llmapi/llm.py
📚 Learning: 2025-08-19T12:45:11.997Z
Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 7033
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:0-0
Timestamp: 2025-08-19T12:45:11.997Z
Learning: In tensorrt_llm/_torch/pyexecutor/model_engine.py, DoRA (Delta Orthogonal Rank Adaptation) functionality was removed from the PyTorch flow to eliminate issues with inverted DoRA detection logic. The original is_dora condition was checking if scaling_vec_pointer == 0, which was potentially incorrect.
Applied to files:
tensorrt_llm/_torch/pyexecutor/model_engine.py
📚 Learning: 2025-08-19T12:45:35.429Z
Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 7033
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:2086-2092
Timestamp: 2025-08-19T12:45:35.429Z
Learning: DoRA (Delta Orthogonal Rank Adaptation) functionality has been removed from the PyTorch flow in tensorrt_llm/_torch/pyexecutor/model_engine.py. The is_dora field is computed but not used downstream in the PyTorch flow, so converting it to a tensor would be wasteful overhead.
Applied to files:
tensorrt_llm/_torch/pyexecutor/model_engine.py
📚 Learning: 2025-08-26T06:07:02.166Z
Learnt from: shaharmor98
Repo: NVIDIA/TensorRT-LLM PR: 7231
File: tensorrt_llm/_torch/pyexecutor/_util.py:504-509
Timestamp: 2025-08-26T06:07:02.166Z
Learning: In tensorrt_llm/_torch/pyexecutor/_util.py, when calling model_engine.set_lora_model_config(), pass model_binding_config.mlp_hidden_size directly without multiplying by mapping.tp_size, as the mlp_hidden_size from get_bindings_model_config() is already the per-TP rank value needed for LoRA weight packaging.
Applied to files:
tensorrt_llm/_torch/pyexecutor/model_engine.py
🧬 Code graph analysis (6)
tensorrt_llm/disaggregated_params.py (1)
tensorrt_llm/_torch/pyexecutor/llm_request.py (2)
mrope_position_ids_handle(368-371)mrope_position_deltas_handle(374-377)
tensorrt_llm/llmapi/mm_encoder.py (1)
triton_backend/all_models/tests/test_llmapi_python_backend.py (1)
inputs(135-153)
tensorrt_llm/llmapi/llm.py (2)
tensorrt_llm/_torch/pyexecutor/llm_request.py (3)
mrope_position_ids_handle(368-371)mrope_position_deltas_handle(374-377)get(129-147)tensorrt_llm/inputs/multimodal.py (1)
MultimodalParams(197-521)
tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
tensorrt_llm/inputs/multimodal.py (1)
to_device(405-449)
tensorrt_llm/_torch/models/modeling_qwen2vl.py (2)
tensorrt_llm/inputs/registry.py (7)
register_input_processor(544-570)support_multimodal_disaggregated(514-541)config(173-175)config(343-344)tokenizer(167-169)tokenizer(338-339)dtype(179-181)tensorrt_llm/_torch/models/modeling_llava_next.py (4)
get_prompt_token_ids(170-243)config(69-70)tokenizer(73-74)dtype(85-86)
tensorrt_llm/executor/result.py (1)
tensorrt_llm/_torch/pyexecutor/llm_request.py (2)
mrope_position_ids_handle(368-371)mrope_position_deltas_handle(374-377)
🪛 Ruff (0.14.8)
tensorrt_llm/_torch/pyexecutor/sampler.py
288-288: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
tensorrt_llm/llmapi/llm.py
468-470: Avoid specifying long messages outside the exception class
(TRY003)
tensorrt_llm/_torch/models/modeling_qwen2vl.py
1073-1073: Avoid specifying long messages outside the exception class
(TRY003)
1076-1076: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (13)
tensorrt_llm/_torch/pyexecutor/llm_request.py (4)
265-266: LGTM!Instance variables follow the existing pattern for
_mm_embeddingsand are properly initialized toNone.
298-306: LGTM!The
set_mrope_positionmethod correctly wraps tensors inSharedTensorContainerand dumps to dict, following the same pattern asappend_mm_embeddings. The parentheses around theSharedTensorContainercalls are unnecessary but don't affect functionality.
367-377: LGTM!The new properties follow the same pattern as
mm_embedding_handleand include helpful docstring notes explaining how to rebuild theSharedTensorContainerfrom the returned dict.
406-410: LGTM!The
py_result_propertiesfrozenset is correctly updated to include the new mrope handle properties, enabling access throughLlmResult.__getattr__.tensorrt_llm/llmapi/mm_encoder.py (1)
105-107: LGTM!The loop refactor from enumeration to direct iteration is cleaner since the index is not used. This aligns with the removal of the unused
_item_athelper function.tensorrt_llm/llmapi/llm.py (3)
92-98: LGTM!Formatting change only - the
_repr_fieldslist items are now on separate lines for better readability.
427-427: LGTM!Using
getattr(..., False)is safer than direct attribute access, preventingAttributeErrorifsupport_mm_disaggis not defined on the input processor.
444-457: LGTM!The mrope_config construction follows a clear pattern: check if
mrope_position_ids_handleis present, assert thatmrope_position_deltas_handleis also present (since they must come as a pair), then build the config dict. The assertion is appropriate here since this is an internal invariant.tensorrt_llm/_torch/pyexecutor/sampler.py (1)
284-305: [Rewritten review comment]
[Classification tag]tensorrt_llm/_torch/pyexecutor/model_engine.py (1)
1699-1714: CPU-onlyto_device()formrope_position_deltasis a good fix.Guarding the transfer by
device.type == "cpu"avoids calling pin/transfer paths on CUDA tensors (which matches the earlier rationale in past review notes).tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py (2)
21-29: Model-dir based parameterization + constructor updates look good.This removes the model-key indirection and makes the test matrix explicit.
Also applies to: 52-65, 68-73
174-184: Encoder max-batch-size matrix is a pragmatic way to encode model quirks.The explicit
(model_dir, encoder_max_batch_size)pairing is much easier to track than conditional skips/config plumbing.Also applies to: 206-213
tensorrt_llm/_torch/models/modeling_qwen2vl.py (1)
978-978:num_context_requests = attn_metadata.num_contextslooks right for disagg threading.This avoids mixing generation requests into the “context embeddings/mrope” preparation slice.
achartier
left a comment
achartier
left a comment
There was a problem hiding this comment.
LGTM for the pyexecutor part, but needs to address the coderabbit comments.
c770297 to
df2c5b1
Compare
2ez4bz
changed the title
|
/bot run |
|
PR_Github #27990 [ run ] triggered by Bot. Commit: |
|
PR_Github #29188 [ run ] triggered by Bot. Commit: |
|
PR_Github #29188 [ run ] completed with state
|
64d4eee to
93dec2b
Compare
|
/bot run |
|
PR_Github #29205 [ run ] triggered by Bot. Commit: |
|
PR_Github #29205 [ run ] completed with state
|
|
/bot run |
|
PR_Github #29212 [ run ] triggered by Bot. Commit: |
|
PR_Github #29212 [ run ] completed with state
|
|
/bot run |
|
PR_Github #29290 [ run ] triggered by Bot. Commit: |
|
PR_Github #29290 [ run ] completed with state
|
* Why? Certain VLMs like the Qwen family need more than just the multimodal embeddings in the language model, and need MRoPE position IDs and deltas. Prior to this commit, only the embeddings could be communicated from the encoder worker to the prefill worker. * What? This commit extends the `DisaggregatedParams` to include the MRoPE information. It also adjusts several pieces of code required to communicate that between E, P and D workers. Closes TRTLLM-9409. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
93dec2b to
d0a511f
Compare
|
/bot run |
|
PR_Github #29329 [ run ] triggered by Bot. Commit: |
|
PR_Github #29329 [ run ] completed with state
|
|
/bot run |
|
PR_Github #29341 [ run ] triggered by Bot. Commit: |
|
PR_Github #29341 [ run ] completed with state
|
|
/bot run |
|
PR_Github #29383 [ run ] triggered by Bot. Commit: |
|
PR_Github #29383 [ run ] completed with state |
|
|
* Why? Certain VLMs like the Qwen family need more than just the multimodal embeddings in the language model, and need MRoPE position IDs and deltas. Prior to this commit, only the embeddings could be communicated from the encoder worker to the prefill worker. * What? This commit extends the `DisaggregatedParams` to include the MRoPE information. It also adjusts several pieces of code required to communicate that between E, P and D workers. Closes TRTLLM-9409. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
* Why? Certain VLMs like the Qwen family need more than just the multimodal embeddings in the language model, and need MRoPE position IDs and deltas. Prior to this commit, only the embeddings could be communicated from the encoder worker to the prefill worker. * What? This commit extends the `DisaggregatedParams` to include the MRoPE information. It also adjusts several pieces of code required to communicate that between E, P and D workers. Closes TRTLLM-9409. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
* Why? Certain VLMs like the Qwen family need more than just the multimodal embeddings in the language model, and need MRoPE position IDs and deltas. Prior to this commit, only the embeddings could be communicated from the encoder worker to the prefill worker. * What? This commit extends the `DisaggregatedParams` to include the MRoPE information. It also adjusts several pieces of code required to communicate that between E, P and D workers. Closes TRTLLM-9409. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> Signed-off-by: Daniil Kulko <kulkodaniil@gmail.com>
9 participants
Summary by CodeRabbit
New Features
Improvements
✏️ Tip: You can customize this high-level summary in your review settings.
Description
Certain VLMs like the Qwen family need more than just the multimodal
embeddings in the language model, and need MRoPE position IDs and
deltas. Prior to this commit, only the embeddings could be communicated
from the encoder worker to the prefill worker.
This commit extends the
DisaggregatedParamsto include the MRoPEinformation. It also adjusts several pieces of code required to
communicate that between E, P and D workers.
Closes TRTLLM-9409.
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
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
Details
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.