[None][perf] TRTLLM MoE maps to lower tuning buckets when ep>1 #9998
Conversation
rosenrodt
force-pushed
the
autotune-sep-bucket-mapper
branch
from
c5e9635 to
0b07662
Compare
📝 WalkthroughWalkthroughThis PR introduces runtime-specific bucket mapping to the autotuner system by adding a Changes
Sequence DiagramsequenceDiagram
actor Tuner as Tuning Phase
participant Cache as Cache System
participant KeyGen as Key Generator
participant Mapper as Bucket Mapper
actor Runtime as Runtime Phase
rect rgb(200, 240, 255)
note over Tuner,Mapper: Tuning-Time Flow (use_tuning_mapping=True)
Tuner->>KeyGen: get_cache_key(spec, use_tuning_mapping=True)
KeyGen->>Mapper: map_to_tuning_buckets(tensor_spec)
Mapper-->>KeyGen: tuning buckets [1,2,4,8,16,32]
KeyGen->>Cache: Store with tuning bucket key
Cache-->>KeyGen: Cache entry created
end
rect rgb(200, 255, 220)
note over Runtime,Mapper: Runtime-Time Flow (use_tuning_mapping=False)
Runtime->>KeyGen: get_cache_key(spec, use_tuning_mapping=False)
KeyGen->>Mapper: map_to_runtime_buckets(tensor_spec, ep_size)
Mapper-->>KeyGen: runtime mapped buckets (scaled by ep_size)
KeyGen->>Cache: Lookup using mapped bucket key
Cache-->>KeyGen: Retrieve matching tuning entry
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes
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: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py (1)
1-16: Missing NVIDIA copyright header.Per coding guidelines, all TensorRT-LLM Open Source Software code files should contain an NVIDIA copyright header that includes the current year at the top.
Add the NVIDIA copyright header at the top of the file:
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from dataclasses import dataclass, replace
🧹 Nitpick comments (1)
tests/unittest/_torch/misc/test_autotuner.py (1)
223-227: Unused variablerunnershould be prefixed with underscore.The static analysis tool correctly identifies that
runneris unpacked but never used. Prefix it with an underscore to indicate it's intentionally ignored.# Verify cache lookup succeeds with the mapped bucket x = torch.randn(input_size, 64) - runner, tactic = tuner.choose_one("test_runtime_bucket_mapping", + _runner, tactic = tuner.choose_one("test_runtime_bucket_mapping", runners, tuning_config, [x, w])
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
tensorrt_llm/_torch/autotuner.py(8 hunks)tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py(26 hunks)tests/unittest/_torch/misc/test_autotuner.py(1 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:
tests/unittest/_torch/misc/test_autotuner.pytensorrt_llm/_torch/autotuner.pytensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.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:
tests/unittest/_torch/misc/test_autotuner.pytensorrt_llm/_torch/autotuner.pytensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py
🧠 Learnings (10)
📓 Common learnings
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.449Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.
📚 Learning: 2025-08-08T04:10:19.038Z
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6728
File: cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp:966-966
Timestamp: 2025-08-08T04:10:19.038Z
Learning: TensorRT plugins currently don't support padding functionality, and TensorRT is not getting new features (in maintenance mode). This means that duplicating parameters like mExpertHiddenSize in function calls, even with TODO comments, can be acceptable as pragmatic solutions within these constraints.
Applied to files:
tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py
📚 Learning: 2025-11-14T11:22:03.729Z
Learnt from: nzmora-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 9163
File: tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py:107-113
Timestamp: 2025-11-14T11:22:03.729Z
Learning: In TensorRT-LLM AutoDeploy custom ops, when adding hardware capability checks to select between kernel implementations (e.g., cuBLAS vs. CUDA kernel), use descriptive variable names that identify the specific GPU architectures or families being targeted (e.g., `is_blackwell_geforce_or_ada`) rather than generic names like `enable_cuda_core`. This makes it clear that the code is selecting an implementation path based on hardware capabilities, not enabling/disabling hardware features.
Applied to files:
tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py
📚 Learning: 2025-08-14T23:23:27.449Z
Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.449Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.
Applied to files:
tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.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/custom_ops/trtllm_gen_custom_ops.py
📚 Learning: 2025-09-19T21:28:13.751Z
Learnt from: jhaotingc
Repo: NVIDIA/TensorRT-LLM PR: 7856
File: cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp:159-166
Timestamp: 2025-09-19T21:28:13.751Z
Learning: In TensorRT-LLM blockScaleMoe routing (cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu), the DeepSeek routing method performs reinterpret_cast<float*>(routingLogits) at line 89, which could cause issues if routing_logits are BF16. However, Qwen3-FP8 models use RenormalizeNaive routing method and are not affected by this dtype casting issue.
Applied to files:
tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py
📚 Learning: 2025-08-20T07:43:36.447Z
Learnt from: ChristinaZ
Repo: NVIDIA/TensorRT-LLM PR: 7068
File: cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh:169-172
Timestamp: 2025-08-20T07:43:36.447Z
Learning: In TensorRT-LLM MOE kernels, when processing up to 128 experts across 32 threads, each thread handles at most 4 experts (N < 5 constraint), where N represents candidates per thread rather than total system capacity.
Applied to files:
tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.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:
tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.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/_torch/custom_ops/trtllm_gen_custom_ops.py
📚 Learning: 2025-12-12T10:07:31.564Z
Learnt from: lirundong
Repo: NVIDIA/TensorRT-LLM PR: 9725
File: tensorrt_llm/_torch/custom_ops/cuda_tile_custom_ops.py:110-178
Timestamp: 2025-12-12T10:07:31.564Z
Learning: In PyTorch custom operators registered with torch.library.custom_op, mutable operators that return None and specify mutates_args do not require a register_fake decorator. Mutation tracking is handled automatically without needing a FakeTensor kernel. This applies to Python custom op definitions in tensorrt_llm/_torch/custom_ops that use mutates_args and return None; verify you are not relying on register_fake in these cases.
Applied to files:
tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py
🧬 Code graph analysis (2)
tests/unittest/_torch/misc/test_autotuner.py (3)
tensorrt_llm/_torch/autotuner.py (6)
get(594-597)TuningConfig(57-111)DynamicTensorSpec(24-39)autotune(235-266)choose_one(672-842)get_specific_custom_op(412-413)tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py (2)
gen_tuning_buckets(63-67)map_to_tuning_buckets(69-71)tensorrt_llm/_torch/utils.py (1)
get_power_of_2_num_tokens_buckets(265-273)
tensorrt_llm/_torch/autotuner.py (2)
tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py (2)
gen_tuning_buckets(63-67)map_to_tuning_buckets(69-71)tests/unittest/_torch/misc/test_autotuner.py (1)
bucket_mapper(177-178)
🪛 Ruff (0.14.8)
tests/unittest/_torch/misc/test_autotuner.py
224-224: Unpacked variable runner is never used
Prefix it with an underscore or any other dummy variable pattern
(RUF059)
⏰ 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 (11)
tensorrt_llm/_torch/custom_ops/trtllm_gen_custom_ops.py (4)
280-289: Lambda closure correctly capturesep_sizefor runtime bucket mapping.The implementation properly differentiates between tuning-time mapping (using
ep_size_=1to not deflate) and runtime mapping (using the actualep_sizeto deflate inflated buffer sizes back to expected token counts). This aligns well with the PR objective of selecting kernels for average-case rather than worst-case token counts during MoE inference.
210-211: EP size calculation pattern is consistent and correct.The
ep_size = num_experts // local_num_expertscalculation correctly derives the expert parallelism size. This pattern is consistently applied across all MoE runner classes.Ensure that
local_num_expertsis always > 0 when these runners are instantiated, as division by zero would cause a runtime error. Based on the MoE architecture, this should always be the case for valid configurations.
617-636: Consistent ep_size-aware implementation across all MoE runners.The
get_dynamic_tensor_specsimplementation inFP8BlockScaleMoERunnerfollows the same pattern asFP4BlockScaleMoERunner, correctly usingmap_to_tuning_bucketswithep_size_=1andmap_to_runtime_bucketswith the actualep_size.
409-409: Correct use of instance-level tuning_config.Using
kernel_runner.tuning_configinstead of the class-level config ensures the ep_size-aware dynamic tensor specs are properly passed through to the autotuner, enabling differentiated runtime vs tuning bucket mapping.tests/unittest/_torch/misc/test_autotuner.py (2)
156-171: Well-documented test for runtime bucket mapping.The docstring clearly explains the distinction between
map_to_tuning_buckets(used during tuning to store cache keys) andmap_to_runtime_buckets(used during runtime to map input sizes to tuning buckets). The MoE EP use case is well-motivated.
196-203: Good verification of cache key storage semantics.The test correctly verifies that cache entries store raw tuning bucket values (not deflated values), ensuring the tuning infrastructure works as intended.
tensorrt_llm/_torch/autotuner.py (5)
31-39: Clean API extension for runtime bucket mapping.The addition of
map_to_runtime_bucketsas an optional field with clear documentation provides a clean separation of concerns. WhenNone, it falls back tomap_to_tuning_buckets, maintaining backward compatibility.
358-387: Correct default foruse_tuning_mappingin cache search.Using
use_tuning_mapping=Falseas the default ensures that runtime cache lookups usemap_to_runtime_buckets(when provided), which correctly maps inflated buffer sizes to the tuning buckets.
1129-1135: Correct bucket mapper selection logic.The conditional logic correctly handles all cases:
- During tuning (
use_tuning_mapping=True): usesmap_to_tuning_buckets- During runtime with custom mapper: uses
map_to_runtime_buckets- During runtime without custom mapper: falls back to
map_to_tuning_bucketsThis ensures backward compatibility while enabling the new EP-aware bucketing.
785-802: Consistent use ofuse_tuning_mapping=Trueduring profiling.Both cache lookup and cache key generation during the tuning loop correctly use
use_tuning_mapping=True, ensuring cache entries are stored and searched using the same bucket mapping (without EP deflation).
750-751: Runtime cache search correctly uses default mapping.The cache search without explicit
use_tuning_mappingparameter defaults toFalse, which means runtime inputs will be mapped usingmap_to_runtime_buckets(when provided). This enables the EP-aware deflation for MoE inference lookups.
rosenrodt
force-pushed
the
autotune-sep-bucket-mapper
branch
from
0b07662 to
d4b456b
Compare
|
/bot run |
|
PR_Github #28524 [ run ] triggered by Bot. Commit: |
|
PR_Github #28524 [ run ] completed with state |
| - Runtime input 4 -> maps to bucket 1 via round_rule(4) = 4 // 4 = 1 | ||
| - Runtime input 16 -> maps to bucket 4 via round_rule(16) = 16 // 4 = 4 | ||
| In MoE EP, the input buffer size is inflated by factor of the EP size to expect the worse case. |
There was a problem hiding this comment.
I am trying to understand this part. Does this mean:
- The input shape during tuning does not have the same definition as the one in runtime?
- Suppose warm-up tuning input num_tokens = 32. During inference, when we see a input with num_tokens = 32, but the op in fact only has a equal workload 32 // 4 = 8 as the tuning phase. Thus, we must use the bucket 8 instead of 32 to get the optimal kernel for our input 32. This becomes the critical distinct between two scenarios.
There was a problem hiding this comment.
Both points are correct. The distinction is when doing tuning the buffer is assumed 100% full and when inferencing the buffer is expected to be more closer to 1/ep_size full than to 100% full.
There was a problem hiding this comment.
Does this optimization also applies to other MoE backends?
There was a problem hiding this comment.
For now, only TRTLLM MoE. I am not sure yet whether CUTLASS MoE has a similar input buffer token count % to TRTLLM MoE at DEP>1. Maybe @bobboli can comment
There was a problem hiding this comment.
The input tensor shape to MoE module are the same across MoE backends, all like [ep_size * max_num_tokens_per_dp_rank, hidden_size]. However, Cutlass MoE has a separate expandRowInputs kernel to extract valid tokens into a separate buffer, so that the actual input to GroupGemm is not inflated anymore. Therefore, there is no discrepancy between tuning and running. @syuoni may confirm.
rosenrodt
force-pushed
the
autotune-sep-bucket-mapper
branch
2 times, most recently
from
aba5ea6 to
901933d
Compare
| m_values = get_last_power_of_2_num_tokens_buckets(MAX_PROFILE_BUCKET) | ||
| round_rule = lambda x: min(last_positive_power_of_2(x), | ||
|
|
||
| # 1/ep_size is the expected token fill rate |
There was a problem hiding this comment.
Consider that each DP rank has n and assume perfect balance,
If ep_size > top_k, in average after dispatch each EP rank has n * top_k tokens, while the provided buffer contains n * ep_size tokens, should we multiply top_k / ep_size, rather than simply divide by ep_size for correction?
If ep_size <= top_k, each EP rank contains all ranks and the input tensor has no invalid tokens, so I think no correction is needed in this case.
There was a problem hiding this comment.
My understanding is that the A2A buffer only holds the routed tokens before the topk expansion. The token count will be effectively expanded by the router kernel (by constructing a topK-to-1 routing table) inside the MoE op. Isn't multiplying top_k like doing the expansion twice?
During tuning, I don't consider how tokens are routed to other ranks. I think of the issue simply as: how sparse is the problem shape? And for DEP, the expectation is only 1/ep_size of the buffer contains actual tokens assuming perfect balance. I am not sure how topk factor fits in the question of how sparse the problem shape is.
There was a problem hiding this comment.
Oh, that top_k factor does not come from router kernel expansion.
Let me describe in a more detailed way:
- Each DP rank contains
ntokens before dispatch. - The recv buffer for AlltoAll on each rank contains space for
n * ep_sizetokens. - If
ep_size > top_k, each DP rank dispatch one token totop_ktarget ranks, without duplication. Therefore, the recv buffers on all ranks containsep_size * n * top_ktokens, in average each recv buffer containsn * top_ktokens. - Therefore, the input tensor to the MoE OP has shape
n * ep_sizebut actually containsn * top_ktokens.
I think the autotuning considers the shape of the input tensor to the MoE OP, right?
There was a problem hiding this comment.
I think the autotuning considers the shape of the input tensor to the MoE OP, right?
Yes
I see your point now. Let me follow your reasoning. If every token is sent to top_k target ranks, then on average each rank's recv buffer contains n * top_k tokens. Any of those tokens gets expanded only to local experts, which is effectively top-1 in perfect balance. Therefore n * top_k tokens in recv buffer of size n * ep_size is still n * top_k after top-1 routing expansion. Exactly the same tokens_after_expansion/buffer_size ratio as when all n tokens are routed to local experts of itself. I think the workload % is eventually the same one way or the other.
Now, does the MoE autotuner guarantee any particular routing behavior (route tokens only to self, or route tokens as wide as possible). The answer is no. I see two issues for MoE autotuner 1) the expert distribution is not aligned across all ranks causing slight rank-to-rank variations in tuning result 2) as result of 1, the workload (measured as tokens after expansion) is not even across all ranks because each rank only routes to self experts.
cc @hyukn this might be the reason why merging distributed autotuning result fails to perform well in e2e. We choose fastest kernel within rank that's is for sure, but when merging among ranks we might need to merge the "slowest" kernel as a way to compensate for unevenly profiled workload.
There was a problem hiding this comment.
Thanks @bobboli and @rosenrodt. These cluse are very innovative. The Cutlass MoE always uses the uniform distribution across all the ranks. This might be the reason why distributed tuning is more stable for that case. Let me taker a further understanding of the analysis.
rosenrodt
force-pushed
the
autotune-sep-bucket-mapper
branch
from
901933d to
fb30e65
Compare
tensorrt_llm/_torch/autotuner.py
Outdated
| opt_shapes.add( | ||
| spec.map_to_tuning_buckets( | ||
| base_profile.shapes[spec.input_idx][spec.dim_idx].val)) | ||
| base_profile.shapes[spec.input_idx][spec.dim_idx].val) |
There was a problem hiding this comment.
@hyukn As discussed offline, I have removed bucket mapping during tuning. The mapping is only applied at inference time now.
|
/bot run |
|
PR_Github #29089 [ run ] triggered by Bot. Commit: |
|
PR_Github #29089 [ run ] completed with state |
| runners, | ||
| p.get_opt_shapes(), | ||
| tuning_config, | ||
| apply_map_to_tuning_buckets=False, |
There was a problem hiding this comment.
Hi @rosenrodt , For TRTLLM Gen MoE, why should we disable apply_map_to_tuning_buckets when do autotuning? Does this affect other operators?
There was a problem hiding this comment.
After discussion with @hyukn , I understand the case for TRTLLM Gen now. The problem is that
- When do autotuning, the per-rank workload is "full", assuming all tokens activate experts at the local rank
- When do inference, the per-rank workload is approximated by
full_workload/ep_size
Normally, we close this gap by inputs_hook, which modifies the inputs when do autotuning. Specific to your case, you can modify the inputs so that the workload is divided by ep_size.
You may refer to the CuteDSL implementation:
TensorRT-LLM/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
Lines 85 to 96 in 237fd0e
Currently, this PR introduces inconsistency between the autotuning and inference shapes, which is a bit concerning.
There was a problem hiding this comment.
@syuoni provides a better option here. Thanks a lot for the suggestion!
Current process of assembling autotuner cache key is:
- Generating tuning_buckets as a list of profiles
- Generate dummy input tensor according to the shape information in the profiles.
- Apply input_pre_hook on the input tensor, and use inputs after pre_hook as the inputs for the runner.forward.
- Generate cache key with
map_to_tuning_bucketmethod.
By defining input_pre_hook, we always generate the tensors with the shapes corresponding to the correct workloads for runner.forward. And the shapes stored in the cache remain to be the original bucket shapes (before input_pre_hook). This means we can also keep map_to_tuning_bucket to a simple bucket mapping method instead of dividing it with ep_size to adjust the workload model.
This actuall extends the usage of input_pre_hook (originally I only want to use it to manipulate the tensor data), but it trully works. I think we should also revise the docstring to clarify this usage.
There was a problem hiding this comment.
@syuoni I fully agree with the input_pre_hook() approach used in cuteDSL but I think it's not directly applicable to TRTLLM MoE.
Let's look at the two main changes in this PR:
First, map_to_tuning_bucket() should not be applied during tuning and this PR addresses that by applying it only during inference. Do you agree that we should keep this change, @syuoni?
- The autotuner tunes the buckets coming solely from
gen_tuning_bucketswithout involvingmap_to_tuning_bucket(). Themap_to_tuning_bucket()then maps the buckets to cache keys which is not the intended behavior as discussed with @hyukn. - This change should not affect the existing ops.
Second—this is the controversial part—the TRTLLM MoE repurposes map_to_tuning_bucket() to account for workload sparsity in a convenient/confusing way depending on how you look at it. Long story short is routing information is not exposed in TRTLLM MoE interface and would require rewrites to fully adopt CuteDSL's approach. I would suggest we defer your suggested approach to a later PR if that's necessary. @syuoni @hyukn let me know what you think :D
- The approach in CuteDSL MoE is sensible. CuteDSL MoE as a module appears to accept routing information at more granular level (tile_idx_to_group_idx, tile_idx_to_mn_limit, etc.)
- TRTLLM MoE accepts either routing_logits or topk_id/top_weights and computes the routing information internally.
- TRTLLM MoE is not able to adopt the same approach without a lot of rewrites around how TRTLLM MoE interacts with the caller.
|
/bot run |
|
PR_Github #30154 [ run ] triggered by Bot. Commit: |
|
PR_Github #30154 [ run ] completed with state
|
|
/bot run |
|
PR_Github #30167 [ run ] triggered by Bot. Commit: |
|
PR_Github #30167 [ run ] completed with state
|
Signed-off-by: Anthony Chang <[email protected]>
Signed-off-by: Anthony Chang <[email protected]>
rosenrodt
force-pushed
the
autotune-sep-bucket-mapper
branch
from
fb30e65 to
4121b67
Compare
|
/bot run |
|
PR_Github #30173 [ run ] triggered by Bot. Commit: |
|
PR_Github #30173 [ run ] completed with state
|
|
/bot run |
|
PR_Github #30181 [ run ] triggered by Bot. Commit: |
|
PR_Github #30181 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #30338 [ run ] triggered by Bot. Commit: |
|
PR_Github #30338 [ run ] completed with state
|
|
/bot run |
|
PR_Github #30358 [ run ] triggered by Bot. Commit: |
|
PR_Github #30358 [ run ] completed with state |
|
/bot run |
|
PR_Github #30372 [ run ] triggered by Bot. Commit: |
|
PR_Github #30372 [ run ] completed with state
|
Signed-off-by: Anthony Chang <[email protected]>
rosenrodt
force-pushed
the
autotune-sep-bucket-mapper
branch
from
41901d1 to
2a4eb96
Compare
6 participants
Summary by CodeRabbit
Release Notes
Refactor
Tests
✏️ Tip: You can customize this high-level summary in your review settings.
Description
The PR allows the autotuner to have separate mappers (one for tuning, one for inference if specified) specifying how token count is mapped to kernel selection buckets.
For MoE inference, the input buffer can hold
runtime_max_tokens_per_rank * ep_sizeto accommodate the worst case scenario where all tokens are routed to single EP rank. The autotuner selects the kernel as if expecting that worst case scenario, but in normal circumstances we should be expecting the average case where the actual filled tokens is closer toruntime_max_tokens_per_rank, essentially undoing theep_sizefactor coming from expecting the worst.Test result with DS-R1
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.