[None][feat] AutoDeploy: Add Nemotron MOE support for AutoDeploy - #8469
Conversation
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run |
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
📝 WalkthroughWalkthroughThis pull request introduces Mixture-of-Experts (MOE) and gated RMSNorm support to the TensorRT LLM auto-deploy framework. Changes include new CUDA and Triton kernels for MOE token alignment, grouped attention pattern splitting by repeat_kv variants, gated RMSNorm operators, NemotronH MOE forward patches, and enhanced KV-cache size tracking. Configuration transforms are updated accordingly. Changes
Sequence Diagram(s)sequenceDiagram
participant App as Application
participant MoE as torch_moe
participant PackTokens as _pack_routed_tokens
participant Kernels as CUDA Kernels
participant Experts as Per-Expert MLPs
App->>MoE: torch_moe(hidden, topk_ids, topk_weights, w1, w2, w3)
MoE->>MoE: Route tokens by top-k experts
MoE->>PackTokens: Pack expert routing
PackTokens->>Kernels: moe_align_block_size (CUDA)
Kernels-->>PackTokens: Grouped token indices, padding
PackTokens-->>MoE: Routed token layout
MoE->>Experts: Dispatch to per-expert MLP
alt mlp_style = "gated_mlp"
Experts->>Experts: y = W2(act(W1·x) ⊙ W3·x)
else mlp_style = "mlp"
Experts->>Experts: y = W_down(act(W_up·x))
end
Experts-->>MoE: Expert outputs
MoE->>MoE: Aggregate with routing weights
MoE-->>App: Combined output
sequenceDiagram
participant Input
participant Triton as Triton Kernel
participant GEMM1 as GEMM (W_up)
participant Act as Activation
participant GEMM2 as GEMM (W_down)
participant Agg as Aggregation
participant Output
Input->>Triton: x, topk_ids, topk_weights
Triton->>GEMM1: hidden = x @ W_up (per-expert)
GEMM1-->>Act: intermediate
Act->>Act: apply_activation(intermediate)
Act-->>GEMM2: activated
Triton->>GEMM2: output = activated @ W_down (per-expert)
GEMM2-->>Agg: per-expert outputs
Agg->>Agg: Apply routing weights
Agg->>Agg: Mask invalid tokens
Agg-->>Output: Final MoE output
sequenceDiagram
participant Norm as torch_rmsnorm_gated
participant Input as x, weight, gate
participant Reshape as Reshape 2D
participant Compute as _layer_norm_fwd
participant Gate as Optional Gate
participant Output as Output
Input->>Reshape: Original shape stored
Reshape->>Reshape: Reshape to (batch, -1)
Reshape->>Compute: Call _layer_norm_fwd(is_rms_norm=True)
Compute->>Compute: Compute RMS normalization
Compute-->>Gate: normalized output
alt gate is not None
Gate->>Gate: norm_before_gate? apply gate before/after norm
Gate-->>Output: gated output
else gate is None
Gate-->>Output: normalized output
end
Output->>Output: Reshape to original shape
Output-->>Norm: Return result
Estimated code review effort🎯 4 (Complex) | ⏱️ ~70 minutes This change introduces substantial new functionality spanning multiple domains: a full CUDA kernel for MOE token alignment (requiring careful memory management and warp-level synchronization review), Triton-based fused MoE operations, new custom operators with tracing support, split transform patterns requiring logic verification, and model-specific patches. The heterogeneity of changes—mixing kernel code, PyTorch custom ops, graph transforms, and configuration—demands separate reasoning for each area. Critical aspects include correctness of CUDA shared-memory layout, Triton kernel semantics, proper routing weight application, and activation resolver dispatch logic. Existing operator modifications (RMSNorm gating, MOE MLP styles) and test coverage for new paths add moderate additional review burden. Pre-merge checks and finishing touches❌ Failed checks (2 warnings)
✅ 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: 15
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (5)
tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py (1)
1-1: Missing NVIDIA Apache-2.0 header.Add the 2025 NVIDIA Apache-2.0 header at the top per coding guidelines.
tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py (1)
1-1: Missing NVIDIA Apache-2.0 header.Add the 2025 NVIDIA Apache-2.0 header at the top per coding guidelines.
tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py (3)
321-326: Possible AttributeError when BFS finds nothing.Check for None before accessing
.args[0].- selected_experts = bfs( + one_hot_node = bfs( common_ancessor2, lambda node: is_op(node, torch.ops.aten.one_hot), attr_next="all_input_nodes", boundary=start_boundary, - ).args[0] - if not selected_experts: + ) + if one_hot_node is None or not one_hot_node.args: continue + selected_experts = one_hot_node.args[0]
93-101: Use typing.List/Optional for Python 3.8 compatibility.Annotations like
list[Node]require 3.9+. Replace withList[Node].-def _find_lowest_common_ancessor(nodes: list[Node]) -> Optional[Node]: +from typing import List, Optional + +def _find_lowest_common_ancessor(nodes: List[Node]) -> Optional[Node]: ... -def _find_final_hidden_state_node( - pattern_output_nodes: list[Node], end_boundary: Node +def _find_final_hidden_state_node( + pattern_output_nodes: List[Node], end_boundary: Node ) -> Optional[Node]: ... -def _extract_index_branches_from_expert_outputs( - pattern_output_nodes: list[Node], -) -> tuple[list[Node], list[Node]]: +def _extract_index_branches_from_expert_outputs( + pattern_output_nodes: List[Node], +) -> Tuple[List[Node], List[Node]]:Also applies to: 290-296, 341-370
1-1: Missing NVIDIA Apache-2.0 header.Add the 2025 NVIDIA Apache-2.0 header at the top per coding guidelines.
🧹 Nitpick comments (23)
tensorrt_llm/_torch/auto_deploy/custom_ops/load_moe_align.py (2)
11-13: Make arch selection adaptive (avoid slow multi-arch compiles by default).Defaulting to multiple SMs increases build time; also misses newer arch (e.g., 10.0). Prefer deriving from the current device if env unset.
-os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "8.0;8.6;8.9;9.0") +if "TORCH_CUDA_ARCH_LIST" not in os.environ and torch.cuda.is_available(): + major, minor = torch.cuda.get_device_capability() + os.environ["TORCH_CUDA_ARCH_LIST"] = f"{major}.{minor}"
15-23: Harden build-directory fallback.Catching only PermissionError can miss Read-only FS or other OSErrors; also consider per-user isolation to avoid collisions.
-try: - os.makedirs(BUILD_DIR, exist_ok=True) -except PermissionError: - import tempfile - - BUILD_DIR = os.path.join(tempfile.gettempdir(), "_moe_align_build") - os.makedirs(BUILD_DIR, exist_ok=True) +try: + os.makedirs(BUILD_DIR, exist_ok=True) +except OSError: + import tempfile, getpass + user = os.environ.get("USER") or os.environ.get("USERNAME") or "unknown" + BUILD_DIR = os.path.join(tempfile.gettempdir(), f"_moe_align_build_{user}") + os.makedirs(BUILD_DIR, exist_ok=True)tensorrt_llm/_torch/auto_deploy/custom_ops/moe_align_kernel.cu (2)
36-42: Reduce redundant shared-memory zeroing.Every thread redundantly zeros the same locations. Zero once per warp or stride over the array.
- for (int i = 0; i < experts_per_warp; ++i) - { - if (my_expert_start + i < padded_num_experts) - { - shared_counts[warp_id * experts_per_warp + i] = 0; - } - } + if ((threadIdx.x % WARP_SIZE) == 0) { + for (int i = 0; i < experts_per_warp; ++i) { + if (my_expert_start + i < padded_num_experts) { + shared_counts[warp_id * experts_per_warp + i] = 0; + } + } + }
10-12: Prefer constexpr/inline over macros for constants.Replace CEILDIV/WARP_SIZE macros with constexpr or inline functions for type-safety and readability.
-#define CEILDIV(x, y) (((x) + (y) -1) / (y)) -#define WARP_SIZE 32 +constexpr int WARP_SIZE = 32; +template <typename T> +__host__ __device__ inline T CEILDIV(T x, T y) { return (x + y - 1) / y; }tensorrt_llm/_torch/auto_deploy/transform/library/attention.py (3)
455-457: Type hint Optional using PEP 604 syntax.Ruff RUF013: prefer bool | None over implicit Optional.
-def generate_and_register_grouped_attn_patterns( - patterns, register_ad_pattern: Callable, only_repeat_kv: bool = None -): +def generate_and_register_grouped_attn_patterns( + patterns, register_ad_pattern: Callable, only_repeat_kv: bool | None = None +):Also applies to: 467-473
542-573: Silence unused parameters in _apply.Prefix with “_” to satisfy ARG002 without changing signature.
- def _apply( - self, - gm: GraphModule, - cm: CachedSequenceInterface, - factory: ModelFactory, - shared_config: SharedConfig, - ) -> Tuple[GraphModule, TransformInfo]: + def _apply( + self, + gm: GraphModule, + _cm: CachedSequenceInterface, + _factory: ModelFactory, + _shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]:Apply similarly in MatchGroupedAttentionWithoutRepeatKV._apply.
Also applies to: 575-616
603-607: Minor: tighten warning text.Nit: “Group Attention Pattern” → “Grouped Attention pattern” for consistency with keys.
- "Fail to find any Group Attention Pattern (without repeat_kv), " + "Failed to find any Grouped Attention pattern (without repeat_kv); "tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_mamba_rms_norm.py (1)
4-4: Remove unused noqa.Ruff RUF100: the “# noqa: F401” isn’t needed here. Delete it.
-import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 +import tensorrt_llm._torch.auto_deploy.custom_opstensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py (3)
11-11: Avoid importing refs from model patches (coupling + cycle risk).Move
_rms_norm_refto a neutral utils/ref module and import from there to prevent transform↔model-patch coupling and potential import cycles.
206-211: Use CUDA dummies (consistency with other patterns).Other RMSNorm patterns trace with CUDA tensors; align dummy args to device="cuda" to avoid topology drift from dtype/device conversions during tracing.
- x = torch.randn(B, S, H, dtype=torch.float32) - w = torch.randn(H, dtype=torch.float32) - g = torch.randn(B, S, H, dtype=torch.float32) + x = torch.randn(B, S, H, device="cuda", dtype=torch.float32) + w = torch.randn(H, device="cuda", dtype=torch.float32) + g = torch.randn(B, S, H, device="cuda", dtype=torch.float32)
185-199: Silence unused params to satisfy linters without changing API.Prefix unused
_applyparams with_to address ARG002 while keeping signature shape.- def _apply( - self, - gm: GraphModule, - cm: CachedSequenceInterface, - factory: ModelFactory, - shared_config: SharedConfig, - ) -> Tuple[GraphModule, TransformInfo]: + def _apply( + self, + gm: GraphModule, + _cm: CachedSequenceInterface, + _factory: ModelFactory, + _shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]:tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_nemotron_h_patches.py (3)
60-63: Skip on CPU-only CI.Guard the test when CUDA isn’t available.
-@pytest.mark.parametrize("B,S", [(2, 6), (1, 8)]) +@pytest.mark.parametrize("B,S", [(2, 6), (1, 8)]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for this test")
41-44: Silence unused loop var.Rename to
_to satisfy linters.- for name, mod in model.named_modules(): + for _, mod in model.named_modules():
46-48: Keep error concise per TRY003.Shorten message (or add details in assertion text if needed).
- if nemotron_moe is None: - raise RuntimeError("NemotronHMOE layer not found. Check your model id or config.") + if nemotron_moe is None: + raise RuntimeError("NemotronHMOE layer not found")tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py (4)
18-21: Iterating while mutating FX graph is fragile.Convert to a list snapshot before edits to avoid iterator invalidation when erasing/inserting nodes.
- for node in graph.nodes: + for node in list(graph.nodes): if not is_op(node, torch.ops.auto_deploy.torch_moe): continue ... - node.replace_all_uses_with(new_node) - graph.erase_node(node) + node.replace_all_uses_with(new_node) + graph.erase_node(node)Also applies to: 68-89
35-48: Make zip strict to catch expert-list length mismatches.This will surface silent mismatches early.
- fused_w_up_experts = torch.stack( - [ - torch.cat( - [gm.get_parameter(w3_node.target), gm.get_parameter(w1_node.target)], dim=-2 - ) - for w1_node, w3_node in zip(w1_list, w3_list) - ], - dim=0, - ) + fused_w_up_experts = torch.stack( + [ + torch.cat( + [gm.get_parameter(w3_node.target), gm.get_parameter(w1_node.target)], + dim=-2, + ) + for w1_node, w3_node in zip(w1_list, w3_list, strict=True) + ], + dim=0, + )
57-58: Normalize mlp_style for robustness.Guard against case differences.
- elif mlp_style_val == "mlp": + elif str(mlp_style_val).lower() == "mlp":
383-392: Erase-node helper should be None-safe.
graph.erase_node(None)may raise the wrong exception; check first.- try: - node_to_remove = bfs(start_boundary, target, attr_next="users", boundary=end_boundary) - graph.erase_node(node_to_remove) - return True - except RuntimeError: - return False + node_to_remove = bfs(start_boundary, target, attr_next="users", boundary=end_boundary) + if node_to_remove is None: + return False + graph.erase_node(node_to_remove) + return Truetests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/triton_kernels/test_triton_moe.py (1)
4-4: Remove unused noqa and keep side-effect import.Switch to a dynamic import to satisfy Ruff (RUF100) while preserving op registration side effects.
-import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 +import importlib +importlib.import_module("tensorrt_llm._torch.auto_deploy.custom_ops")tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py (2)
86-141: Clarify dtype contract (return fp32 vs input dtype).This op returns fp32, unlike torch_rmsnorm which preserves input dtype. If downstream expects input dtype, cast back; otherwise document the fp32 contract prominently.
Option A (preserve input dtype):
- return out2.reshape(x_shape) + return out2.reshape(x_shape).to(x.dtype)Option B: keep fp32 but add explicit note where this is consumed. Please confirm intended behavior.
143-151: Silence unused-arg lints in fake/meta.Ruff flags eps and norm_before_gate as unused. Discard them explicitly.
def _torch_rmsnorm_gated_meta( x, weight, gate, eps: float, group_size: int, norm_before_gate: bool = False, ): + # unused in meta path + del eps, norm_before_gatetensorrt_llm/_torch/auto_deploy/custom_ops/triton_moe.py (2)
218-225: Tighten ImportError message for extension load.Message is long and over-specific (TRY003). Keep it concise and surface original exception.
- except Exception as e: - raise ImportError( - f"Failed to load moe_align_block_size CUDA extension. " - f"Error: {e}. Make sure CUDA toolkit and nvcc are available." - ) from e + except Exception as e: + raise ImportError("Failed to load moe_align_block_size extension") from e
236-248: Drop unused argument or use it.num_tokens_post_padded is passed but unused in _invoke_kernel. Either remove it from signature and call sites, or use it to bound EM/grid.
-def _invoke_kernel( +def _invoke_kernel( A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, topk_weights: torch.Tensor | None, sorted_token_ids: torch.Tensor, expert_ids: torch.Tensor, - num_tokens_post_padded: torch.Tensor, # Changed to tensor for CUDA graph compatibility mul_routed_weight: bool, top_k: int, config: dict, compute_type, ):And drop the corresponding actual arguments at both call sites.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (17)
tensorrt_llm/_torch/auto_deploy/config/default.yaml(2 hunks)tensorrt_llm/_torch/auto_deploy/custom_ops/load_moe_align.py(1 hunks)tensorrt_llm/_torch/auto_deploy/custom_ops/moe_align_kernel.cu(1 hunks)tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py(2 hunks)tensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.py(3 hunks)tensorrt_llm/_torch/auto_deploy/custom_ops/triton_moe.py(1 hunks)tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py(3 hunks)tensorrt_llm/_torch/auto_deploy/shim/interface.py(1 hunks)tensorrt_llm/_torch/auto_deploy/transform/library/attention.py(5 hunks)tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py(1 hunks)tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py(2 hunks)tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py(2 hunks)tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_mamba_rms_norm.py(1 hunks)tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/triton_kernels/test_triton_moe.py(1 hunks)tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_nemotron_h_patches.py(1 hunks)tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_attention_matcher.py(1 hunks)tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_attention_matcher_hf.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Use only spaces, no tabs; indent with 4 spaces.
Files:
tensorrt_llm/_torch/auto_deploy/shim/interface.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_attention_matcher_hf.pytensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.pytensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/triton_kernels/test_triton_moe.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_attention_matcher.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_mamba_rms_norm.pytensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.pytensorrt_llm/_torch/auto_deploy/custom_ops/triton_moe.pytensorrt_llm/_torch/auto_deploy/transform/library/kvcache.pytests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_nemotron_h_patches.pytensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.pytensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.pytensorrt_llm/_torch/auto_deploy/transform/library/attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/load_moe_align.pytensorrt_llm/_torch/auto_deploy/custom_ops/moe_align_kernel.cu
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py: Python code must target Python 3.8+.
Indent Python code with 4 spaces; do not use tabs.
Maintain module namespace when importing; prefer 'from package.subpackage import foo' then 'foo.SomeClass()' instead of importing the class directly.
Python filenames should be snake_case (e.g., some_file.py).
Python classes use PascalCase names.
Functions and methods use snake_case names.
Local variables use snake_case; prefix 'k' for variables that start with a number (e.g., k_99th_percentile).
Global variables use upper SNAKE_CASE prefixed with 'G' (e.g., G_MY_GLOBAL).
Constants use upper SNAKE_CASE (e.g., MY_CONSTANT).
Avoid shadowing variables from an outer scope.
Initialize all externally visible members of a class in the constructor.
Prefer docstrings for interfaces that may be used outside a file; comments for in-function or file-local interfaces.
Use Google-style docstrings for classes and functions (Sphinx-parsable).
Document attributes and variables inline so they render under the class/function docstring.
Avoid reflection when a simpler, explicit approach suffices (e.g., avoid dict(**locals()) patterns).
In try/except, catch the most specific exceptions possible.
For duck-typing try/except, keep the try body minimal and use else for the main logic.
Files:
tensorrt_llm/_torch/auto_deploy/shim/interface.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_attention_matcher_hf.pytensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.pytensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/triton_kernels/test_triton_moe.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_attention_matcher.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_mamba_rms_norm.pytensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.pytensorrt_llm/_torch/auto_deploy/custom_ops/triton_moe.pytensorrt_llm/_torch/auto_deploy/transform/library/kvcache.pytests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_nemotron_h_patches.pytensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.pytensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.pytensorrt_llm/_torch/auto_deploy/transform/library/attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/load_moe_align.py
**/*.{cpp,cxx,cc,h,hpp,hh,hxx,cu,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Prepend the NVIDIA Apache-2.0 copyright header with current year to the top of all source files (e.g., .cpp, .h, .cu, .py).
Files:
tensorrt_llm/_torch/auto_deploy/shim/interface.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_attention_matcher_hf.pytensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.pytensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/triton_kernels/test_triton_moe.pytests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_attention_matcher.pytests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_mamba_rms_norm.pytensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.pytensorrt_llm/_torch/auto_deploy/custom_ops/triton_moe.pytensorrt_llm/_torch/auto_deploy/transform/library/kvcache.pytests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_nemotron_h_patches.pytensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.pytensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.pytensorrt_llm/_torch/auto_deploy/transform/library/attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/load_moe_align.pytensorrt_llm/_torch/auto_deploy/custom_ops/moe_align_kernel.cu
**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh}: Namespace closing braces must include a trailing comment with the namespace name (e.g., '} // namespace foo').
Prefer const or constexpr variables over #define for constants.
Declare variables that are not modified after initialization as const.
Avoid magic literals in code; except for 0, nullptr, true, false. Use named constants for comparisons and logic.
Use Allman brace style for formatting.
Place the semicolon of an empty for/while loop on a new line.
Bodies of switch/while/do-while/for must be compound statements (brace-delimited), and if/else must always be followed by brace-delimited statements.
Type names (e.g., classes) must be CamelCase starting with an uppercase letter (e.g., FooBar).
Local variables, methods, and namespaces use lowerCamelCase (e.g., localFooBar).
Non-magic-number global variables that are non-static and not in an anonymous namespace must be lowerCamelCase prefixed with 'g' (e.g., gDontUseGlobalFoos).
Non-magic-number globals that are static or in an anonymous namespace use lowerCamelCase prefixed with 's' (e.g., sMutableStaticGlobal).
Locally visible static variables use lowerCamelCase with 's' prefix (e.g., static std::once_flag sFlag).
Private/protected member variables use 'm' prefix with CamelCase (e.g., mNbFooValues). Public members may omit, but 'm' is encouraged for clarity.
Constants (enums, global constants, static constants, and function-scope magic/literal constants) use uppercase SNAKE_CASE with 'k' prefix (e.g., kDIGIT_NUM).
Function-scope constants that are not magic numbers or literals are named like non-constant variables (e.g., bool const pass = a && b).
If macros are necessary, name them in UPPER_SNAKE_CASE (e.g., FOO_VERSION) and prefer constants over #define.
Use LLVM clang-format; wrap lines at a maximum of 120 columns; use '// clang-format off/on' sparingly with justification.
Use smart pointers for heap allocations; prefer unique_ptr for sole ownership, shared_ptr for shared...
Files:
tensorrt_llm/_torch/auto_deploy/custom_ops/moe_align_kernel.cu
**/*.{cpp,cxx,cc,cu,h,hpp,hh,hxx,cuh}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
C++ filenames should be lowerCamelCase (first letter lowercase) and must be case-insensitive unique within a compilation target.
Files:
tensorrt_llm/_torch/auto_deploy/custom_ops/moe_align_kernel.cu
🧬 Code graph analysis (11)
tensorrt_llm/_torch/auto_deploy/shim/interface.py (1)
tensorrt_llm/_utils.py (1)
numel(1022-1023)
tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py (3)
tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py (1)
_rms_norm_ref(21-44)tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py (1)
torch_rmsnorm_gated(87-140)tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py (2)
ADPatternMatcherPass(61-67)register_ad_pattern(99-182)
tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py (3)
tensorrt_llm/_torch/auto_deploy/utils/node_utils.py (2)
extract_op_args(407-444)is_op(179-202)tensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.py (1)
torch_moe(66-135)tensorrt_llm/module.py (1)
register_parameter(186-190)
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/triton_kernels/test_triton_moe.py (2)
tensorrt_llm/_torch/auto_deploy/custom_ops/load_moe_align.py (1)
moe_align_block_size(43-67)tensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.py (1)
torch_moe(66-135)
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_mamba_rms_norm.py (2)
tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py (1)
_rms_norm_ref(21-44)tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py (1)
torch_rmsnorm_gated(87-140)
tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py (1)
tensorrt_llm/_torch/modules/fla/layernorm_gated.py (1)
_layer_norm_fwd(117-180)
tensorrt_llm/_torch/auto_deploy/custom_ops/triton_moe.py (1)
tensorrt_llm/_torch/auto_deploy/custom_ops/load_moe_align.py (1)
moe_align_block_size(43-67)
tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py (2)
tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py (2)
num_pages(321-322)num_pages(325-329)tensorrt_llm/_torch/auto_deploy/transform/interface.py (1)
_log_info(420-422)
tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_nemotron_h_patches.py (2)
tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py (1)
_nemotron_h_moe_forward(92-116)tensorrt_llm/_torch/model_config.py (1)
torch_dtype(204-209)
tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py (1)
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.py (1)
torch_moe(66-135)
tensorrt_llm/_torch/auto_deploy/transform/library/attention.py (2)
tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py (2)
register_ad_pattern(99-182)ADPatternMatcherPass(61-67)tensorrt_llm/_torch/auto_deploy/transform/interface.py (3)
TransformRegistry(503-531)register(509-516)BaseTransform(213-500)
🪛 Ruff (0.14.0)
tensorrt_llm/_torch/auto_deploy/transform/library/rms_norm.py
195-195: Unused method argument: cm
(ARG002)
196-196: Unused method argument: factory
(ARG002)
197-197: Unused method argument: shared_config
(ARG002)
tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py
41-41: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
57-57: Avoid specifying long messages outside the exception class
(TRY003)
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/triton_kernels/test_triton_moe.py
4-4: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_mamba_rms_norm.py
4-4: Unused noqa directive (non-enabled: F401)
Remove unused noqa directive
(RUF100)
tensorrt_llm/_torch/auto_deploy/custom_ops/rms_norm.py
148-148: Unused function argument: eps
(ARG001)
150-150: Unused function argument: norm_before_gate
(ARG001)
tensorrt_llm/_torch/auto_deploy/custom_ops/triton_moe.py
157-157: Unused function argument: N
(ARG001)
157-157: Unused function argument: K
(ARG001)
157-157: Unused function argument: top_k
(ARG001)
221-224: Avoid specifying long messages outside the exception class
(TRY003)
243-243: Unused function argument: num_tokens_post_padded
(ARG001)
346-346: Avoid specifying long messages outside the exception class
(TRY003)
430-430: Unused function argument: selected_experts
(ARG001)
431-431: Unused function argument: routing_weights
(ARG001)
432-432: Unused function argument: w1_stacked_weight
(ARG001)
433-433: Unused function argument: w2_stacked_weight
(ARG001)
tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_nemotron_h_patches.py
41-41: Loop control variable name not used within loop body
Rename unused name to _name
(B007)
47-47: Avoid specifying long messages outside the exception class
(TRY003)
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.py
26-26: Avoid specifying long messages outside the exception class
(TRY003)
133-133: Avoid specifying long messages outside the exception class
(TRY003)
146-146: Unused function argument: mlp_style
(ARG001)
147-147: Unused function argument: act_fn
(ARG001)
tensorrt_llm/_torch/auto_deploy/transform/library/attention.py
456-456: PEP 484 prohibits implicit Optional
Convert to T | None
(RUF013)
588-588: Unused method argument: cm
(ARG002)
589-589: Unused method argument: factory
(ARG002)
590-590: Unused method argument: shared_config
(ARG002)
⏰ 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 (8)
tensorrt_llm/_torch/auto_deploy/shim/interface.py (1)
70-79: LGTM! KV-specific cache sizing accessor looks good.The implementation correctly filters KV caches by prefix and calculates total memory usage. The docstring clearly explains the exclusion of non-scaling caches.
tensorrt_llm/_torch/auto_deploy/transform/library/kvcache.py (3)
255-258: Good backward compatibility pattern.The defensive accessor with fallback ensures compatibility with older versions of
CachedSequenceInterfacethat may not have the KV-specific sizing method.
260-267: Nice logging improvements.Removing the incorrect "(MB)" unit from the page count and adding conditional KV-only logging both improve clarity without adding noise.
291-300: Review edge case: per-page calculation when KV cache size is zero.The calculation at Line 299 sets
per_page_bytesto 1 whencurrent_kv_cache_sizeis 0, which would causenew_num_pagesto equalnew_kv_total_bytes(potentially billions of pages if measured in bytes).While this transform appears to run after
initialize_cache(making zero KV cache size unlikely), confirm that:
- KV caches are always initialized before resize, or
- Add a guard to handle/skip the resize when
current_kv_cache_size == 0tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_attention_matcher_hf.py (1)
44-49: Ordering looks correct.Running “with_repeat_kv” before “without_repeat_kv” prevents over-matching and aligns with the transform docstring.
If possible, add a tiny assertion that the “with_repeat_kv” pass finds ≥1 pattern in at least one parametrized case to guard against regressions.
tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py (1)
82-84: Treating "moe" like "mlp" is reasonable here.Forward path reuse looks fine.
tensorrt_llm/_torch/auto_deploy/config/default.yaml (1)
36-39: Config additions look consistent.Split grouped-attention matchers and new fuse_gated_rmsnorm entry align with code changes. Ensure downstream docs/tests reference these exact keys.
Also applies to: 128-130
tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_attention_matcher.py (1)
449-454: Config split verified and correct; no issues found.The configuration changes have been successfully verified:
- Both
match_grouped_attention_with_repeat_kvandmatch_grouped_attention_without_repeat_kvare registered in default.yaml (lines 36–38)- The old bare
match_grouped_attentionentry has been removed- Test file (lines 449–454) aligns with default.yaml configuration
- Order is preserved:
with_repeat_kvexecutes beforewithout_repeat_kv, which is semantically correct (more specific case first, fallback second)
|
/bot run |
|
PR_Github #21773 [ run ] triggered by Bot. Commit: |
|
PR_Github #21773 [ run ] completed with state |
Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run |
|
PR_Github #21940 [ run ] triggered by Bot. Commit: |
|
PR_Github #21940 [ run ] completed with state |
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run |
|
/bot run |
|
PR_Github #21951 [ run ] triggered by Bot. Commit: |
|
PR_Github #21951 [ run ] completed with state |
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Chenghao Zhang <211069071+nvchenghaoz@users.noreply.github.com>
|
/bot run |
|
PR_Github #21972 [ run ] triggered by Bot. Commit: |
|
PR_Github #21972 [ run ] completed with state |
|
/bot run |
|
PR_Github #22073 [ run ] triggered by Bot. Commit: |
|
PR_Github #22073 [ run ] completed with state |
|
/bot run |
|
PR_Github #22085 [ run ] triggered by Bot. Commit: |
|
@nvchenghaoz, this is a large change. Can you explain what new dependencies were added or modified to simplify the OSS review process? |
|
PR_Github #22085 [ run ] completed with state |
atrifex
left a comment
atrifex
left a comment
There was a problem hiding this comment.
Conditional on no new dependencies being added.
|
The .cu file is added for triton moe BF16 support which gives the speed up during the end-to-end model run. |
…DIA#8469) Signed-off-by: yufeiwu-nv <230315618+yufeiwu-nv@users.noreply.github.com>
6 participants
Summary by CodeRabbit
Release Notes
New Features
Tests
This is a mass merge to support the Nemotron MOE structure. The key change includes: