[TRTLLM-5966][feat] Initial steps towards Helix parallelism support #5668
Conversation
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #10612 [ run ] triggered by Bot |
|
PR_Github #10612 [ run ] completed with state |
|
/bot run |
|
PR_Github #10651 [ run ] triggered by Bot |
|
PR_Github #10651 [ run ] completed with state |
|
/bot run |
|
PR_Github #10655 [ run ] triggered by Bot |
|
PR_Github #10655 [ run ] completed with state |
|
/bot run |
|
PR_Github #10662 [ run ] triggered by Bot |
|
PR_Github #10662 [ run ] completed with state |
|
/bot run |
|
PR_Github #10788 [ run ] triggered by Bot |
|
/bot run |
|
PR_Github #10803 [ run ] triggered by Bot |
|
PR_Github #10788 [ run ] completed with state |
|
PR_Github #10803 [ run ] completed with state |
|
/bot run |
|
PR_Github #10832 [ run ] triggered by Bot |
|
PR_Github #10832 [ run ] completed with state |
|
sorry if anyone has started looking at this: it should have been a draft PR from the start but I only noticed now |
Signed-off-by: Matthias Jouanneaux <[email protected]>
|
PR_Github #12870 [ run ] completed with state |
Signed-off-by: Matthias Jouanneaux <[email protected]>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
tensorrt_llm/_torch/distributed/ops.py (1)
337-338: Add missing None value handling for tensor lists.The function doesn't handle
Nonevalues in tensor lists like the existingallgatherfunction does. This could cause issues when processing lists containingNonevalues.Consider adding the same
Nonehandling logic used inallgather:else: + input, valid = filter_valid_input(input) torch_op = torch.ops.trtllm.allgather_list output_info = [get_output_info(val, dim) for val in input] input = [ val.contiguous().view(-1, val_info['numel_base']) for val, val_info in zip(input, output_info) ]And restore the full output at the end:
else: output = [ convert_output(val, val_info) for val, val_info in zip(output, output_info) ] + output = restore_full_output(output, valid) return output
🧹 Nitpick comments (2)
tensorrt_llm/_torch/attention_backend/trtllm.py (1)
219-219: Fix line length to comply with style guidelines.The documentation is clear and well-formatted, but the line exceeds the 120-character limit.
- helix_position_offsets (torch.Tensor): The tensor to store the helix position offsets, with shape (num_tokens) on GPU. + helix_position_offsets (torch.Tensor): The tensor to store the helix position offsets, + with shape (num_tokens) on GPU.tensorrt_llm/_torch/distributed/ops.py (1)
244-244: Fix line length violation.The line exceeds the 120-character limit specified in coding guidelines.
Apply this diff to fix the line length:
- If None (default), the splits are concatenated along dimension given by `dims`. + If None (default), the splits are concatenated along dimension given by + `dims`.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
tensorrt_llm/_torch/attention_backend/trtllm.py(6 hunks)tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py(1 hunks)tensorrt_llm/_torch/distributed/ops.py(1 hunks)tensorrt_llm/_torch/pyexecutor/model_engine.py(2 hunks)tensorrt_llm/_torch/pyexecutor/py_executor.py(2 hunks)tensorrt_llm/_torch/pyexecutor/resource_manager.py(2 hunks)tests/integration/test_lists/waives.txt(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- tests/integration/test_lists/waives.txt
🚧 Files skipped from review as they are similar to previous changes (4)
- tensorrt_llm/_torch/pyexecutor/model_engine.py
- tensorrt_llm/_torch/pyexecutor/resource_manager.py
- tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
- tensorrt_llm/_torch/pyexecutor/py_executor.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.py: Python code 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.
Python filenames should use snake_case (e.g., some_file.py).
Python classes should use PascalCase (e.g., class SomeClass).
Python functions and methods should use snake_case (e.g., def my_awesome_function():).
Python local variables should use snake_case. Prefix k for variable names that start with a number (e.g., k_99th_percentile).
Python global variables should use upper snake_case and prefix G (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 interfaces that may be used outside a Python file, prefer docstrings over comments.
Comments in Python 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.
Attributes and variables in Python can be documented inline; attribute docstrings will be rendered under the docstring for the class.
Avoid using reflection in Python when functionality can be easily achieved without it.
When using try-except blocks in Python, limit the except to the smallest set of errors possible.
When using try-except blocks to handle multiple possible variable types in Python, keep the body of the try as small as possible, using the else block to implement the logic.
Files:
tensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/distributed/ops.py
**/*.{cpp,h,hpp,cc,cxx,cu,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.
Files:
tensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/distributed/ops.py
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/attention_backend/trtllm.py
219-219: Line too long (130 > 120)
(E501)
tensorrt_llm/_torch/distributed/ops.py
244-244: Line too long (131 > 120)
(E501)
🔇 Additional comments (7)
tensorrt_llm/_torch/attention_backend/trtllm.py (5)
178-178: LGTM! Parameter addition follows established patterns.The
helix_position_offsetsparameter is properly typed and positioned appropriately in the method signature.
256-256: LGTM! Proper instance attribute assignment.The assignment follows the established pattern for storing tensor parameters as instance attributes.
469-474: LGTM! Good refactoring to group tensor parameters.The refactoring to group related tensor parameters into a single list is a good practice that avoids parameter limits while maintaining clarity. The addition of
helix_position_offsetsto this grouped list is appropriate.
1096-1096: LGTM! Consistent parameter addition across API.The
helix_position_offsetsparameter addition maintains consistency with theplanmethod and follows proper typing patterns.
1163-1163: LGTM! Proper parameter propagation completed.The
helix_position_offsetsparameter is correctly passed to the wrapper's plan method, completing the proper data flow through the call stack.tensorrt_llm/_torch/distributed/ops.py (2)
225-300: Well-implemented all-to-all collective operation.The function correctly implements all-to-all communication with proper:
- Parameter validation and type checking
- Even splitting of input tensors across ranks
- Integration with NCCL-based C++ backend
- Output reconstruction with flexible concatenation/stacking options
The logic for splitting inputs and reconstructing outputs is sound and handles both single tensors and tensor lists appropriately.
302-385: CP-specific all-gather implementation looks correct.The function properly implements all-gather for the context parallelism group with:
- Correct use of
mapping.cp_groupandmapping.cp_size- Proper size optimization logic (lines 340-344)
- Consistent tensor reshaping and output reconstruction logic
The implementation follows the same pattern as the existing
allgatherfunction, which is good for maintainability.
|
/bot run |
|
PR_Github #12984 [ run ] triggered by Bot |
|
PR_Github #12984 [ run ] completed with state |
|
/bot run |
|
PR_Github #13005 [ run ] triggered by Bot |
|
PR_Github #13005 [ run ] completed with state |
|
/bot run |
|
PR_Github #13013 [ run ] triggered by Bot |
|
PR_Github #13013 [ run ] completed with state |
Signed-off-by: Matthias Jouanneaux <[email protected]>
coderabbitai
bot
added
the
Community want to contribute
|
/bot run |
|
PR_Github #13154 [ run ] triggered by Bot |
|
PR_Github #13154 [ run ] completed with state |
|
/bot run |
|
PR_Github #13212 [ run ] triggered by Bot |
|
PR_Github #13212 [ run ] completed with state |
|
/bot run |
|
PR_Github #13224 [ run ] triggered by Bot |
|
PR_Github #13224 [ run ] completed with state |
There was a problem hiding this comment.
Hi Matthias, thank you for the awesome work!
I'm half-way through the review and am worried the current change is too big and is doing too many things. I think it's better to break it down into 4-5 MRs for more focused reviews. My proposal would be to merge in this order:
- Add alltoall op to TRTLLM
- Re-group attentionOp args to deal with len<64 requirement
- Changes to Mapping class
- Core changes to MLA kernels and attention.py
- Changes for disagg
Do you think that's a good idea?
| common_enqueue_params.workspace = workspace_ptr; | ||
| if (softmax_stats_tensor.has_value()) | ||
| { | ||
| TLLM_CHECK_WITH_INFO(softmax_stats_tensor.value().size(-1) == 2, |
There was a problem hiding this comment.
Why 2? Can we please mention the expected values in cpp/tensorrt_llm/common/attentionOp.h where this is defined?
There was a problem hiding this comment.
it's for both the max and LSE (denominator) of softmax. Since I'm not adding the softmax stats variable to the underlying op in this PR, I'll try to see where it is defined and whether this is already mentioned
| torch::optional<torch::Tensor> mrope_rotary_cos_sin, torch::optional<torch::Tensor> mrope_position_deltas, | ||
| std::optional<torch::Tensor> mla_context_paged_kv, std::optional<torch::Tensor> mla_context_kv_cache_block_offsets, | ||
| std::optional<int64_t> attention_chunk_size, std::optional<torch::Tensor> softmax_stats_tensor, | ||
| std::optional<int64_t> attention_chunk_size, std::vector<torch::optional<torch::Tensor>> additional_tensors, |
There was a problem hiding this comment.
The limitation of 64 args is unfortunate. Is this grouping logical though? They seem a bit disjoint.
Also, will be nice to leave a comment explaining what each index in additional_tensors is.
There was a problem hiding this comment.
The grouping isn't logical but I'm not seeing a logical way of how the arguments are organized. I think recent PRs also ran into this limit and reduced the number of arguments, so maybe we can just add it without issues now
| // optional when cross attention | ||
| int32_t const* encoder_input_lengths = nullptr; | ||
| int64_t const* runtime_perf_knobs = nullptr; | ||
| // optional when compute attention stats (MLA chunked prefill or Helix parallelism) |
There was a problem hiding this comment.
Do you think it's better to maintain two separate variables - one for chunked prefill and another for helix parallelism?
What if both chunked prefill and helix parallelism are at play at the same time? Is sharing the same variable still the right thing to do?
There was a problem hiding this comment.
no, both chunked prefill and Helix use the exact same logical output from attention. It is not shared because chunked prefill and Helix will not call attention at the same time.
| def __init__(self, mapping: Mapping): | ||
| super().__init__(mapping) | ||
| self.create_tp_comm() | ||
| self.create_cp_comm() |
There was a problem hiding this comment.
Dumb question: How come this was not needed by previously existing types of CP? Did they not need any NCCL collectives?
There was a problem hiding this comment.
the torch backend basically doesn't support any type of CP. The only thing it supports is star attention, which is completely custom (own backend etc.) and AFAIK, star attention cannot be mixed with any of the other parallelisms, or at least only in very specific ways
| pp_size = config.mapping.pp_size * config.mapping.tp_size // overridden_tp_size | ||
| mapping = Mapping( | ||
| world_size=tp_size * pp_size, | ||
| world_size=tp_size * pp_size * self.mapping.cp_size, |
There was a problem hiding this comment.
The paper differentiates tp_size for attention (TPA) and FFN (TPF). But, we tend to share tp_size throughout to mean both so far. From this change, looks like we're essentially reusing tp_size to mean TPA (please correct me if I'm wrong).
Using cp_size in GatedMLP seems out of place to me. Maybe we should move this upstream and send in overriden_tp_size instead?
There was a problem hiding this comment.
we can move this to another place indeed by just passing a different mapping to GatedMLP .
| if moe_cluster_size == -1: | ||
| moe_cluster_size = 1 | ||
|
|
||
| cp_type = CpType.ULYSSES if cp_config is None else cp_config.get( |
There was a problem hiding this comment.
Why default to Ulysses of cp_config is None? Should we have a CpType.NONE instead?
There was a problem hiding this comment.
This was not added by me, previously Ulysses was the de-facto default (although it was not explicitly mentioned). It is also only for non-torch backend (as mentioned elsewhere, torch backend doesn't support any CP right now).
I have no idea who decided that Ulysses should be the default.
|
|
||
| def has_cp_helix(self): | ||
| return self.cp_size > 1 and self.cp_config.get( | ||
| "cp_type", CpType.HELIX) == CpType.HELIX |
There was a problem hiding this comment.
If cp_config doesn't have cp_type, has_cp_ulysses() and has_cp_helix() will both return True which doesn't sound right to me. Is this intentional?
Full disclosure: I don't know what Ulysses does. XD
There was a problem hiding this comment.
Indeed that seems like an issue, I'll fix it.
|
@MatthiasKohl @brb-nv can we close this PR? |
4 participants
[TRTLLM-5966][feat] Initial steps towards Helix parallelism support
See JIRA ticket https://jirasw.nvidia.com/browse/TRTLLM-5966
Description
This PR adds initial steps for TRT-LLM to be able to support Helix parallelism:
all_rank_num_tokenstoall_tp_rank_num_tokens/all_cp_rank_num_tokens. Note: after merging latest main, this is not done throughout the code-base for the AttentionMetadata anymore.Test Coverage
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 [--disable-fail-fast --skip-test --stage-list "A10-1, xxx" --gpu-type "A30, H100_PCIe" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-[Post-Merge]-1, xxx"]Launch build/test pipelines. All previously running jobs will be killed.
--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-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-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.--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. Will also run 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-[Post-Merge]-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-[Post-Merge]-1, xxx".For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.md.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.
Summary by CodeRabbit
New Features
Improvements
Bug Fixes
Chores