[None][feat] AutoDeploy: propagate layer_type hint across pattern-matcher rewrites by greg-kwasniewski1 · Pull Request #14835 · NVIDIA/TensorRT-LLM

greg-kwasniewski1 · 2026-06-01T21:10:22Z

Summary

Generic, centralized propagation of the layer-level layer_type sharding hint across AutoDeploy pattern-matcher rewrites, so it survives N->1 fusions (e.g. the SwiGLU matchers) and the downstream hint-driven sharder (apply_sharding_hints / shard_layers) still sees it.
Root-causes and generalizes the Qwen3.5-MoE NVFP4 shared-expert accuracy regression that [https://nvbugs/6221450][fix] AutoDeploy: Qwen3.5 400B NVFP4 accuracy regression fix #14667 fixed point-wise.

Background

The SwiGLU matchers collapse three fine-grained linears into one torch_(nvfp4_)swiglu_mlp op before sharding. That fused op carried no layer_type, so the fail-open shard_layers whitelist TP-sharded the shared expert that was meant to be replicated. The hint must survive the rewrite.

Changes

Add an optional layer_type arg to the intermediate SwiGLU ops (torch_swiglu_mlp, torch_nvfp4_swiglu_mlp, torch_finegrained_fp8_swiglu_mlp) and their fakes so a fused node can carry it.
Add collect_classification_hints (consensus across matched nodes; conflicting values dropped with a warning) and stamp_hints (schema-guarded set) in node_utils.py.
Propagate generically in ADReplacementPatternEntry.apply: collect from match.nodes before the rewrite, stamp the inserted op(s) afterward.
Per-weight mechanics (tp_mode, output_sizes, ...) are intentionally NOT propagated -- a fused op's ShardableNode re-derives them structurally and they legitimately differ across constituents (an MLA layer mixes none/colwise/rowwise yet shares one layer_type), so generic propagation of the sharding decision is ill-posed.

Test plan

GPU-free unit tests for the helpers: consensus over mixed tp_mode, conflict-drop, schema-guarded stamping.
NVFP4 end-to-end regression: layer_type survives the SwiGLU match onto the fused op (mirrors the Qwen3.5-MoE bug).
Full tests/unittest/auto_deploy/singlegpu/transformations/library/ suite: 329 passed locally.
CI green (triggered after the PR opens).

Risk

library-visible (matcher infra + custom ops; no public API change).

Notes

Follow-up: once [https://nvbugs/6221450][fix] AutoDeploy: Qwen3.5 400B NVFP4 accuracy regression fix #14667 merges, its per-transform wmap propagation can be removed in favor of this central hook.

Fixes #14834

Summary by CodeRabbit

New Features
- Enhanced graph optimization with layer type classification hints during auto-deployment, enabling improved operation fusion and sharding strategies.
Tests
- Added tests validating classification hint propagation through graph transformations.

greg-kwasniewski1 · 2026-06-01T21:11:24Z

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1"

coderabbitai · 2026-06-01T21:15:09Z

📝 Walkthrough

Walkthrough

This PR implements layer-type hint propagation through AutoDeploy's pattern-matcher infrastructure. SwiGLU ops now carry layer_type metadata, reusable hint-collection utilities ensure consensus across matched nodes, and pattern rewrites preserve those hints on fused ops. Tests validate the infrastructure and demonstrate the fix for Qwen3.5-MoE shared-expert TP-sharding bugs.

Changes

Layer-type hint propagation

Layer / File(s)	Summary
Extend SwiGLU custom ops with layer_type parameter `tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py`	Three quantization variants of `torch_swiglu_mlp` (`torch_nvfp4_swiglu_mlp`, `torch_finegrained_fp8_swiglu_mlp`) now accept optional `layer_type: str = "unknown"` parameter in both their `@torch.library.custom_op` signatures and `register_fake` tracing implementations. Parameter is documented as a sharding classification hint that does not affect numeric results.
Implement hint collection and stamping utilities `tensorrt_llm/_torch/auto_deploy/utils/node_utils.py`	Adds `CLASSIFICATION_HINT_NAMES` constant (`{"layer_type"}`), `collect_classification_hints()` to compute consensus values across matched nodes (ignoring `None`/`"unknown"`, dropping conflicts with warning), and `stamp_hints()` to apply hints only to ops whose op schemas declare the hint arguments.
Integrate hint propagation into pattern matching `tensorrt_llm/_torch/auto_deploy/utils/pattern_matcher.py`	`ADReplacementPatternEntry.apply` now collects classification hints from matched nodes before graph replacement, snapshots the pre-rewrite node set when hints are present, and stamps the newly inserted nodes with those hints after `replace_with_graph`.
Validate hint propagation behavior `tests/unittest/auto_deploy/singlegpu/transformations/library/test_fuse_swiglu.py`, `tests/unittest/auto_deploy/singlegpu/transformations/library/test_nvfp4_swiglu.py`	Three graph-only unit tests verify consensus recovery for shared `layer_type` across mixed mechanics, conflict detection/dropping, and schema-guarded stamping (skipping incompatible ops). Regression test confirms `layer_type="shared_expert"` survives NVFP4 SwiGLU pattern matching onto the fused op.

Sequence Diagram

sequenceDiagram
    participant PatternMatcher
    participant collect_classification_hints as collect hints
    participant replace_with_graph as Graph replace
    participant stamp_hints as stamp hints
    participant FXOp as Fused FX Op
    PatternMatcher->>collect_classification_hints: matched.nodes
    collect_classification_hints-->>PatternMatcher: hints dict
    PatternMatcher->>replace_with_graph: Execute replacement
    replace_with_graph-->>PatternMatcher: new nodes inserted
    PatternMatcher->>stamp_hints: new_nodes + hints
    stamp_hints->>FXOp: Set layer_type arg
    stamp_hints-->>PatternMatcher: count updated

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title clearly describes the main change: propagating layer_type hints across pattern-matcher rewrites in AutoDeploy, which is the primary objective of the changeset.
Linked Issues check	✅ Passed	All code changes align with issue `#14834` objectives: layer_type args added to SwiGLU ops, helper functions implemented in node_utils.py, hint propagation added to pattern_matcher.py, and comprehensive tests validate the behavior.
Out of Scope Changes check	✅ Passed	All changes are directly scoped to implementing the feature: SwiGLU custom ops, helper utilities, pattern matcher integration, and targeted tests—no extraneous modifications detected.
Docstring Coverage	✅ Passed	Docstring coverage is 94.44% which is sufficient. The required threshold is 80.00%.
Description check	✅ Passed	PR description is detailed and comprehensive, covering summary, background, changes, test plan, risk, and notes. All required template sections are either present or appropriately addressed via detailed explanation.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

tensorrt-cicd · 2026-06-01T21:17:12Z

PR_Github #51420 [ run ] triggered by Bot. Commit: 3ce70e9 Link to invocation

tensorrt-cicd · 2026-06-02T03:08:35Z

PR_Github #51420 [ run ] completed with state SUCCESS. Commit: 3ce70e9
/LLM/main/L0_MergeRequest_PR pipeline #40829 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

greg-kwasniewski1 · 2026-06-03T13:31:09Z

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast --post-merge

tensorrt-cicd · 2026-06-03T13:38:22Z

PR_Github #51848 [ run ] triggered by Bot. Commit: 8981253 Link to invocation

tensorrt-cicd · 2026-06-04T03:12:35Z

PR_Github #51848 [ run ] completed with state FAILURE. Commit: 8981253
/LLM/main/L0_MergeRequest_PR pipeline #41207 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

greg-kwasniewski1 · 2026-06-04T11:03:15Z

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast --post-merge

…rites AutoDeploy pattern matchers that collapse N fine-grained ops into one (e.g. the SwiGLU matchers) dropped the per-op layer_type sharding hint. The hint-driven sharder (apply_sharding_hints / shard_layers) then read layer_type=None on the fused op and, because the whitelist is fail-open, TP-sharded nodes meant to be replicated -- the root cause of the Qwen3.5-MoE NVFP4 shared-expert accuracy regression fixed point-wise in NVIDIA#14667. Generalize the fix into the matcher infra: add an optional layer_type arg to the intermediate SwiGLU ops, add collect_classification_hints / stamp_hints helpers in node_utils, and propagate the consensus layer_type from the matched nodes onto the replacement op in ADReplacementPatternEntry. Per-weight mechanics (tp_mode, output_sizes, ...) are intentionally not propagated: the fused op's ShardableNode re-derives them structurally and they legitimately differ across constituents (an MLA layer mixes none/colwise/rowwise yet shares one layer_type), so generic propagation of the sharding decision is ill-posed. Tested: GPU-free unit tests for the helpers (consensus / conflict-drop / schema-guarded stamping), an NVFP4 end-to-end regression mirroring the Qwen3.5-MoE bug, and the full transformations/library suite (329 passed). Fixes NVIDIA#14834 Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

greg-kwasniewski1 · 2026-06-04T11:09:49Z

/bot run --extra-stage "DGX_B200-4_GPUs-AutoDeploy-1, DGX_H100-4_GPUs-AutoDeploy-1" --disable-fail-fast --post-merge

tensorrt-cicd · 2026-06-04T11:10:07Z

PR_Github #52056 [ run ] triggered by Bot. Commit: 89daace Link to invocation

tensorrt-cicd · 2026-06-04T11:16:10Z

PR_Github #52058 [ run ] triggered by Bot. Commit: 89daace Link to invocation

tensorrt-cicd · 2026-06-04T11:19:48Z

PR_Github #52056 [ run ] completed with state ABORTED. Commit: 89daace

Link to invocation

tensorrt-cicd · 2026-06-05T04:01:54Z

PR_Github #52058 [ run ] completed with state FAILURE. Commit: 89daace
/LLM/main/L0_MergeRequest_PR pipeline #41391 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

greg-kwasniewski1 · 2026-06-07T18:32:28Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-07T18:38:43Z

PR_Github #52595 [ run ] triggered by Bot. Commit: 89daace Link to invocation

tensorrt-cicd · 2026-06-07T20:32:05Z

PR_Github #52595 [ run ] completed with state SUCCESS. Commit: 89daace
/LLM/main/L0_MergeRequest_PR pipeline #41877 completed with status: 'SUCCESS'

CI Report

Link to invocation

taylor-yb-lee · 2026-06-07T21:43:58Z

Tested Qwen3.5 with https://github.com/nv-auto-deploy/TensorRT-LLM/tree/taylor/fix_qwen3.5_acc_new_0607 (Removing redundant patches added in #14667
Confirmed that the Qwen3.5 acc passes.

greg-kwasniewski1 requested a review from a team as a code owner June 1, 2026 21:10

greg-kwasniewski1 requested a review from taylor-yb-lee June 1, 2026 21:10

github-actions Bot assigned greg-kwasniewski1 Jun 1, 2026

greg-kwasniewski1 force-pushed the gk/propagate-layer-type-hint branch from 8981253 to 89daace Compare June 4, 2026 11:08

taylor-yb-lee approved these changes Jun 7, 2026

View reviewed changes

greg-kwasniewski1 merged commit 02f6b2f into NVIDIA:main Jun 8, 2026
9 of 10 checks passed

coderabbitai Bot mentioned this pull request Jun 10, 2026

[None][feat] AutoDeploy: Qwen3.5: Apply whielist based sharding and apply lm_head sharding #15185

Merged

1 task

Uh oh!

Conversation

greg-kwasniewski1 commented Jun 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Background

Changes

Test plan

Risk

Notes

Summary by CodeRabbit

Uh oh!

greg-kwasniewski1 commented Jun 1, 2026

Uh oh!

coderabbitai Bot commented Jun 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram

Estimated code review effort

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

greg-kwasniewski1 commented Jun 3, 2026

Uh oh!

tensorrt-cicd commented Jun 3, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

greg-kwasniewski1 commented Jun 4, 2026

Uh oh!

greg-kwasniewski1 commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

greg-kwasniewski1 commented Jun 7, 2026

Uh oh!

tensorrt-cicd commented Jun 7, 2026

Uh oh!

tensorrt-cicd commented Jun 7, 2026

Uh oh!

taylor-yb-lee commented Jun 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

greg-kwasniewski1 commented Jun 1, 2026 •

edited

Loading

coderabbitai Bot commented Jun 1, 2026 •

edited

Loading