[#10345][perf] Enable multi-stream MOE for super. Also adds multi-stream MLA attn - #11520
Conversation
Speed up torch.export for MOE models by temporarily reducing the number of experts during tracing, then programmatically expanding the FX graph to include all expert weights afterward. Expert nn.ModuleLists are discovered via a TorchDispatchMode probe that intercepts torch_moe-family custom ops during a lightweight forward pass, so the optimization does NOT rely on any particular attribute name. Changes: - export.py: Add _MoeExpertProbe / _find_moe_module_lists for name- independent expert discovery. Add _reduce_moe_experts / _restore_moe_experts / _expand_moe_experts_in_graph helpers and wire them into torch_export_to_gm via num_moe_experts_for_export. - export_to_gm.py: Thread num_moe_experts_for_export through the ExportToGMConfig and transform pipeline. - test_export.py: Add SimpleMoEForExport unit tests (parametrized over expert_attr_name='experts' and 'mlp_bank' to prove name independence) and optional Glm4MoeLiteForCausalLM real-model tests. Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
📝 WalkthroughWalkthroughThis PR introduces SwiGLU MLP fusion, multi-stream execution for MoE and MLA attention, reduced MOE expert tracing for efficient export, and refactored Add+RMSNorm fusion. Configuration updates enable these new transform stages across multiple deployment profiles. Changes
Sequence Diagram(s)sequenceDiagram
participant Model as Source Model
participant Probe as MOE Expert Probe
participant Tracer as Torch Exporter
participant GraphMod as GraphModule
participant Expander as MOE Expander
Model->>Probe: Scan for MOE parameters
Probe->>Probe: Identify expert ModuleLists
Probe-->>Model: Return expert targets
Model->>Model: Reduce expert counts (num_moe_experts_for_export)
Model->>Tracer: Forward with reduced experts
Tracer->>GraphMod: Capture reduced graph
GraphMod-->>Tracer: Return captured module
Tracer->>Model: Restore original expert lists
Tracer->>Expander: Expand MOE weights in graph
Expander->>GraphMod: Create get_attr nodes
Expander->>GraphMod: Register new parameters
Expander-->>Tracer: Return expanded GraphModule
Tracer-->>Model: Final optimized graph
sequenceDiagram
participant GraphMod as FX GraphModule
participant Analyzer as MOE Analyzer
participant StreamMgr as CUDA Stream Manager
participant GraphRewriter as Graph Rewriter
GraphMod->>Analyzer: Locate MOE fusion nodes
Analyzer->>Analyzer: Find merge add node downstream
Analyzer->>Analyzer: Identify routed vs shared inputs
Analyzer->>Analyzer: Traverse shared-expert subgraph
Analyzer-->>GraphRewriter: Return fork point and shared ops
GraphRewriter->>StreamMgr: Allocate auxiliary stream
GraphRewriter->>GraphMod: Insert begin_aux_stream_passthrough
GraphRewriter->>GraphMod: Mark shared expert operations
GraphRewriter->>GraphMod: Insert end_aux_stream_passthrough
GraphRewriter->>StreamMgr: Record synchronization event
GraphRewriter->>GraphMod: Insert wait_aux_stream_passthrough on routed path
GraphRewriter->>GraphMod: Rewire merge add dependencies
GraphRewriter-->>GraphMod: Return transformed module with multi-stream ops
sequenceDiagram
participant GraphMod as GraphModule
participant PatternMatch as Pattern Matcher
participant GraphTransform as Graph Transformer
participant FusionEngine as Fusion Engine
GraphMod->>PatternMatch: Scan for SwiGLU patterns
PatternMatch->>PatternMatch: Detect silu(x @ gate.T) * (x @ up.T) @ down.T
PatternMatch->>GraphTransform: Register pattern matches
GraphTransform->>GraphTransform: Replace with torch_swiglu_mlp ops
GraphTransform-->>GraphMod: Update graph with matched ops
GraphMod->>FusionEngine: Fuse matched SwiGLU ops
FusionEngine->>FusionEngine: Concatenate gate + up weights
FusionEngine->>FusionEngine: Optionally concatenate biases
FusionEngine->>GraphMod: Create fused_swiglu_mlp nodes
FusionEngine->>GraphMod: Rewire dependencies
FusionEngine->>GraphMod: Eliminate dead code
FusionEngine-->>GraphMod: Return fused GraphModule
Estimated code review effort🎯 4 (Complex) | ⏱️ ~70 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 3❌ Failed checks (3 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts (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
🤖 Fix all issues with AI agents
In `@tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py`:
- Around line 129-140: The rare-branch handling of cast_node currently builds a
list from cast_node.users and inserts new_cast before an arbitrary user, which
wastes allocation and can violate topological order; change the insertion to use
the add_out node as the anchor (use graph.inserting_after(add_out)) so new_cast
is placed after its input and before all users, and avoid list(...) by using
next(iter(cast_node.users)) if you still need to reference a user; keep the
replace_all_uses_with(new_cast) and graph.erase_node(cast_node) logic and add
id(cast_node) to erased as before.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py`:
- Around line 166-171: The graph currently hard-codes the device by passing
torch.cuda.current_device() into record_event_passthrough inside the
graph.inserting_before(earliest_q) block, which bakes a transform-time device
into the FX graph; remove the device kwarg so record_event_passthrough uses its
default (runtime resolution, -1) and resolves torch.cuda.current_device() at
call time—update the call that creates rec_node (the graph.call_function
invocation of record_event_passthrough) to omit the device argument.
- Around line 1-12: Add the required NVIDIA Apache-2.0 copyright header (with
the year of latest meaningful modification) at the very top of the file before
the existing module docstring in multi_stream_attn.py; ensure the header uses
the standard NVIDIA header text and license notice so it precedes the docstring
and applies to the entire module.
- Around line 217-218: Remove the unconditional file dump that writes gm.graph
to /tmp (the with open(... "/tmp/after_multi_stream_mla_attn.txt") block) — this
is debug-only code; delete that block from multi_stream_attn.py and replace it
with a debug log call using the existing logger (import ad_logger from
...utils.logger if not already present) e.g. call ad_logger.debug with a
descriptive message and gm.graph (e.g., "After multi_stream_mla_attn
transform:\n%s", gm.graph) so the information is only emitted to logs and not
written to disk.
In `@tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py`:
- Around line 435-436: Remove the unconditional debug dump that writes gm.graph
to /tmp — delete the with open("/tmp/after_multi_stream_moe.txt", "w") ...
f.write(str(gm.graph)) block (or gate it behind proper debug logging) inside the
multi-stream MoE transform where the local variable gm is available; if you need
to keep the dump for local debugging, wrap it with a runtime debug check (e.g.,
if logger.isEnabledFor(logging.DEBUG) or an explicit DEBUG_TRANSFORMS flag) and
surround the file write with try/except to avoid raising on read-only
filesystems, and avoid writing to /tmp or any world-readable location.
In `@tensorrt_llm/_torch/auto_deploy/utils/_graph.py`:
- Around line 74-75: Replace the private attribute access base_overload._schema
with the public torch.library API: call
torch.library.infer_schema(base_overload) (or infer_schema(base_op) if you need
the original callable) to obtain the operator schema, and assign that result to
schema; keep the base_overload selection logic (base_overload = base_op.default
if hasattr(base_op, "default") else base_op) but replace the direct _schema
access with torch.library.infer_schema(...) so you rely on the stable
torch.library schema inference API.
🧹 Nitpick comments (14)
tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py (2)
6-8: Remove unusednoqadirective (static analysis).Ruff reports the blanket
noqaon line 6 is unused. If the import itself is used (it is — in_count_fused_ops), just drop the# noqacomment.Also, the import style pulls in a symbol directly rather than using
from package.subpackage import module. The same applies to lines 11–12 (TransformConfig,FuseAddRMSNorm). Per coding guidelines, preferfrom package.subpackage import moduleand then referencemodule.Symbol.Example for lines 6-8 and 11-12
-from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.flashinfer_fused_add_rms_norm import ( # noqa - flashinfer_fused_add_rms_norm, -) +from tensorrt_llm._torch.auto_deploy.custom_ops.normalization import flashinfer_fused_add_rms_norm as fused_add_rms_norm_modHowever, since
flashinfer_fused_add_rms_normis used as a node-target identity check, importing the function directly may be more pragmatic here. At minimum, remove the stale# noqa.-from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.flashinfer_fused_add_rms_norm import ( # noqa - flashinfer_fused_add_rms_norm, -) +from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.flashinfer_fused_add_rms_norm import ( + flashinfer_fused_add_rms_norm, +)
124-130:_count_fused_opschecks Python function identity, not a registered op.
_count_rms_norm_opsand_count_add_opsboth useis_op(n, torch.ops.…)to match registered ops, but_count_fused_opscomparesn.targetagainst the raw Python functionflashinfer_fused_add_rms_norm. This will only work if the fused op is inserted into the graph as a plaincall_functiontargeting that Python callable (not as atorch.opsdispatch key). If the transform implementation changes to register a propertorch.opsentry, this check will silently stop matching.This is fine if the current transform indeed inserts a plain Python callable, but worth keeping in mind.
tensorrt_llm/_torch/auto_deploy/utils/_graph.py (1)
105-110: Default fake implementation assumes output shape equals first arg's shape.The
_default_fakereturnstorch.empty_like(args[0]), which is only correct for ops whose output shape matches the first input. For ops like linear (where output shape differs), callers must always providemake_fake. The docstring documents this, but consider adding a brief inline comment here as well since this is a common footgun.tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py (3)
32-32: Import style: import the module, not the function directly.Per coding guidelines, use
from ...custom_ops.linear import swigluand referenceswiglu.torch_swiglu_mlp. Same applies to line 289 fortorch_nvfp4_swiglu_mlp.Proposed fix
-from ...custom_ops.linear.swiglu import torch_swiglu_mlp +from ...custom_ops.linear import swigluThen update references:
- return torch_swiglu_mlp(x, gate_weight, up_weight, down_weight, None, None, None) + return swiglu.torch_swiglu_mlp(x, gate_weight, up_weight, down_weight, None, None, None)And for line 289:
-from ...custom_ops.linear.swiglu import torch_nvfp4_swiglu_mlp # noqa: E402 +# (move to top-of-file imports) +# from ...custom_ops.linear import swiglu # already imported aboveThen use
swiglu.torch_nvfp4_swiglu_mlpin the replacement function.As per coding guidelines, "Python imports must use
from package.subpackage import modulestyle; never usefrom module import Class."
287-289: Move this import to the top of the file.The mid-file import with
# noqa: E402is unnecessary — there's no circular dependency issue since both symbols come from the same module. Moving it to line 32 alongside the otherswigluimport would be cleaner.
548-563: No validation that gate and up input scales match before fusing.The docstring says "gate+up input_scale and alpha must match (shared input)" but there's no assertion. If they differ (e.g., from a misconfigured quantization pipeline), this would silently use only gate's scales, producing incorrect results.
Proposed assertion
+ # Validate that gate and up share the same input scale and alpha + # (required because the fused op uses a single scale for both). + gate_input_scale = get_attr_by_name(gm, gate_input_scale_node.target) + up_input_scale = get_attr_by_name(gm, up_input_scale_node.target) + gate_alpha = get_attr_by_name(gm, gate_alpha_node.target) + up_alpha = get_attr_by_name(gm, up_alpha_node.target) + assert torch.equal(gate_input_scale, up_input_scale), ( + "Gate and up input scales must match for NVFP4 SwiGLU fusion" + ) + assert torch.equal(gate_alpha, up_alpha), ( + "Gate and up alpha must match for NVFP4 SwiGLU fusion" + ) + # Create the fused_nvfp4_swiglu_mlp nodetensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py (1)
238-240: BFS useslist.pop(0)— O(n) per pop.
queue.pop(0)on a list is O(n). Usecollections.dequewithpopleft()for O(1) BFS. Same issue at line 323.Proposed fix
+from collections import deque + def _find_merge_add(moe_node: Node) -> Optional[Node]: visited: Set[Node] = set() - queue = list(moe_node.users.keys()) + queue = deque(moe_node.users.keys()) while queue: - n = queue.pop(0) + n = queue.popleft() if n in visited: continue visited.add(n) if is_op(n, torch.ops.aten.add.Tensor): return n queue.extend(n.users.keys())Similarly for line 320-323:
- queue = [shared_output] + queue = deque([shared_output]) while queue: - n = queue.pop(0) + n = queue.popleft()tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py (2)
27-32: ImportModelFactorydirectly from its source module.
ModelFactoryis imported transitively throughmulti_stream_moe, but it originates from...models.factory. Import it directly to avoid a hidden dependency chain and make the dependency explicit.Proposed fix
from .multi_stream_moe import ( - ModelFactory, _make_aux_stream_impl, cuda_stream_manager, record_event_passthrough, ) +from ...models.factory import ModelFactory
48-110: KV projection detection heuristic is reasonable but fragile.The "has downstream linear within 3 hops = Q-like, otherwise KV-like" heuristic works well for DeepSeek-style MLA but could misfire on architectures where the KV chain includes a learned transformation (linear). A brief inline comment noting which architecture(s) this targets would help future maintainers.
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_attn.py (1)
19-23: Import style: importing private functions directly rather than the module.The coding guidelines specify
from package.subpackage import modulestyle. Here, individual underscore-prefixed functions are imported directly. Since these are internal test helpers being tested, this is a common pattern in unit tests, but flagging for awareness.tensorrt_llm/_torch/auto_deploy/export/export.py (1)
236-329:_expand_moe_experts_in_graph: hardcoded arg index 3 assumes a fixed MOE op signature.Line 272 starts scanning list arguments from
range(3, len(node.args)). This assumes alltorch_moe-family ops have the same argument layout with per-expert weight lists starting at position 3. If any of the three MOE ops (torch_moe,torch_quant_fp8_moe,torch_quant_nvfp4_moe) has a different signature, this would silently skip or mis-identify weight lists. A defensive check or a per-op signature map would be more robust, though the current code works correctly given the existing op signatures.tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/test_export.py (2)
331-335: Duplicate helper:_count_moe_expertsand_count_moe_experts_in_graphare identical.Both functions iterate the graph looking for
torch_moenodes and returnlen(node.args[3]). Consider extracting a single shared helper.Proposed consolidation
Move
_count_moe_experts_in_graphto module level (outside theif _HAS_GLM4:block) and reuse it intest_moe_export_with_reduced_experts:+def _count_moe_experts_in_graph(gm: GraphModule) -> int: + """Return the number of experts in the first ``torch_moe`` call in *gm*.""" + for node in gm.graph.nodes: + if node.op == "call_function" and "torch_moe" in str(node.target): + return len(node.args[3]) # w1_weight list length + return 0 + + `@pytest.mark.parametrize`("expert_attr_name", ["experts", "mlp_bank"]) ... def test_moe_export_with_reduced_experts(...): ... - def _count_moe_experts(gm): - for node in gm.graph.nodes: - if node.op == "call_function" and "torch_moe" in str(node.target): - return len(node.args[3]) - return 0 - - assert _count_moe_experts(gm_full) == num_experts + assert _count_moe_experts_in_graph(gm_full) == num_expertsAlso applies to: 400-405
249-249: Unusednoqadirective per static analysis.Ruff reports
RUF100: Unused noqa directive (non-enabled: F401, E402). IfF401andE402are not enabled in the project's Ruff config, thenoqacomment is unnecessary.tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py (1)
14-14: Remove unusednoqadirective.Static analysis (Ruff RUF100) flags
# noqa: F401as unnecessary here. The side-effect import is valid, but the suppression comment isn't needed.Proposed fix
-import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 +import tensorrt_llm._torch.auto_deploy.custom_ops
taylor-yb-lee
left a comment
taylor-yb-lee
left a comment
There was a problem hiding this comment.
Looks like need rebase to remove the commit for export
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
- Replace hardcoded magic number 3 with schema-driven introspection for finding List[Tensor] arguments in MOE ops - Rename min_num_experts to num_moe_experts_for_export for consistency - Remove unnecessary try/except guard around GLM4 MoE Lite import in tests Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #35973 [ run ] triggered by Bot. Commit: |
… multi_stream_attn Remove the device kwarg from the graph.call_function invocation of record_event_passthrough so that the default (-1) is used and torch.cuda.current_device() is resolved at runtime rather than being baked into the FX graph at transform time. Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
|
PR_Github #35973 [ run ] completed with state
|
|
/bot run |
|
PR_Github #35975 [ run ] triggered by Bot. Commit: |
|
PR_Github #35975 [ run ] completed with state
|
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #35986 [ run ] triggered by Bot. Commit: |
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #35988 [ run ] triggered by Bot. Commit: |
|
PR_Github #35988 [ run ] completed with state
|
Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com>
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #36004 [ run ] triggered by Bot. Commit: |
|
PR_Github #36011 [ run ] triggered by Bot. Commit: |
|
PR_Github #36011 [ run ] completed with state
|
|
/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" |
|
PR_Github #36018 [ run ] triggered by Bot. Commit: |
|
PR_Github #36018 [ run ] completed with state |
| return _derived_op_libs[namespace] | ||
|
|
||
|
|
||
| def create_derived_custom_op( |
There was a problem hiding this comment.
given that this transform is at a later stage, where we are already outside the AutoDeploy IR, I'd prefer to just use regular call_function nodes calling into eager python rather than defining custom ops.
Might be something good to discuss
There was a problem hiding this comment.
The aux_stream_wrapper (now removed) was an attempt at doing precisely this. It ran into issues with transforms downstream of multi-stream
…ti-stream MLA attn (NVIDIA#11520) Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Signed-off-by: peihu-nv <259410613+peihu-nv@users.noreply.github.com>
4 participants
Summary by CodeRabbit
Release Notes
New Features
Tests