[feat] Sharding logic split to pattern detection and executor for EP and BMM (fixes #5916)

[greg-kwasniewski1]

[https://github.com//issues/5916][feat] EP and BMM Sharding logic split to pattern detection and executor

Description

A follow-up of PR nv-auto-deploy#86 , pattern matcher + executor extended to EP and BMM.

Test Coverage

tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py (renamed previous test_graph_sharding.py)
tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py
tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_bmm_sharding.py

Summary by CodeRabbit

• New Features

  • Introduced a modular, staged graph transformation and optimization pipeline for model deployment, with configurable YAML-based settings and deep merging support.
  • Added support for quantized Mixture-of-Experts (MoE) models (FP8, FP4/NVFP4) and corresponding transformation, fusion, and operator implementations.
  • Added a pure PyTorch attention backend with caching, sliding window, logit capping, and sink token support.
  • Enhanced configuration management with dynamic YAML merging and advanced CLI overrides.
  • Provided detailed documentation and expert configuration guides for AutoDeploy.

• Improvements

  • Refactored sharding, quantization, and pattern matching logic for greater modularity, testability, and extensibility.
  • Updated graph transformation functions to operate in-place, improving performance and clarity.
  • Expanded test coverage for quantized MoE, attention, sharding, and configuration utilities.
  • Improved error handling, metadata tracking, and logging throughout the transformation pipeline.

• Bug Fixes

  • Fixed parameter handling and default merging in model and tokenizer configuration.
  • Addressed issues with input constraint cleanup and aliased parameter loading in graph export.

• Chores

  • Updated dependencies and packaging to support new features and configuration files.
  • Improved and extended documentation for end-users and contributors.

[coderabbitai]

Caution

Review failed

The pull request is closed.

Walkthrough

This update introduces a modular, extensible transformation pipeline for the Torch AutoDeploy backend, adding new configuration management utilities, a staged transform interface, and a library of graph transformations. The refactor replaces direct in-place graph mutations with a staged, metadata-driven approach, deepens YAML-based config merging, and enhances quantization, sharding, and MoE support. Tests are updated to match the new APIs and in-place semantics.

Changes

File(s) / Area	Change Summary
tensorrt_llm/_torch/auto_deploy/transform/, .../library/, .../interface.py, .../optimizer.py	Introduced a modular graph transformation pipeline: new staged transform interface, registry, and library of transforms (build, export, cleanup, quantize, etc.).
tensorrt_llm/_torch/auto_deploy/llm_args.py	Refactored config classes: added AutoDeployConfig (simple) and LlmArgs (advanced), improved validation, and separated config for modular pipeline.
tensorrt_llm/_torch/auto_deploy/utils/_config.py	Added dynamic YAML config mix-in and deep merge utilities for flexible, multi-file configuration.
tensorrt_llm/_torch/auto_deploy/transformations/library/, .../fusion.py, .../quantization.py, .../eliminate_redundant_transposes.py, .../fused_moe.py, .../kvcache.py, .../sharding.py, .../collectives.py, .../attention.py, .../rope.py	Updated all legacy transformation functions to operate in-place (no return), refactored sharding and quantization logic, and added quantized MoE support.
tensorrt_llm/_torch/auto_deploy/transformations/library/ep_sharding.py	Removed expert parallel sharding module (now handled in new staged pipeline).
tensorrt_llm/_torch/auto_deploy/transformations/transform.py	Refactored optimizer to use modular staged pipeline; updated to new config and transform APIs.
tensorrt_llm/_torch/auto_deploy/transformations/export.py, .../_graph.py	Removed input constraint cleanup, improved robustness in aliased parameter handling, and updated canonicalization to in-place.
tensorrt_llm/_torch/auto_deploy/transformations/library/quantize_moe.py	New: Implements quantization passes for MoE nodes (FP8, FP4).
tensorrt_llm/_torch/auto_deploy/models/hf.py	Centralized model/tokenizer defaults, improved merging, and checkpoint loading.
tensorrt_llm/_torch/auto_deploy/build_and_run_ad.py	Enhanced CLI/config handling: deep merging, dynamic YAML, and flexible overrides.
tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py	Updated to use AutoDeployConfig, propagate max_beam_width, and improved input prep.
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_moe.py	Refactored MoE logic, added quantized FP8/FP4 MoE ops, and modularized expert dispatch.
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py	New: Pure torch backend attention operator with caching, sliding window, sinks, and logit capping.
tensorrt_llm/_torch/auto_deploy/custom_ops/triton_attention.py, .../_triton_attention_internal.py, .../triton_kernels/attention_with_kv_cache.py	Added support for sliding window and sinks in attention kernels and interfaces.
tensorrt_llm/_torch/auto_deploy/custom_ops/torch_attention.py	Extended SDPA op to accept sinks, sliding window, and logit cap.
tensorrt_llm/_torch/auto_deploy/utils/quantization_utils.py	Added quantization pattern exclusion, scale extraction, and quantization type detection utilities.
tensorrt_llm/_torch/auto_deploy/utils/node_utils.py	Improved node filtering utilities and docstrings, removed unused helpers.
tensorrt_llm/_torch/auto_deploy/transformations/__init__.py	Added deprecation docstring, pointing to new transform pipeline.
tensorrt_llm/_torch/auto_deploy/transform/__init__.py	New: Entry point for modular transform pipeline, re-exporting interface.
tensorrt_llm/_torch/auto_deploy/transform/library/build_model.py, .../export_to_gm.py, .../cleanup_input_constraints.py	New: Modular transforms for building models, exporting to graph module, and cleaning input constraints.
tensorrt_llm/_torch/auto_deploy/transform/optimizer.py	New: Modular inference optimizer applying staged transforms by config.
tensorrt_llm/_torch/auto_deploy/transform/interface.py	New: Transform interface, registry, config, and metadata management.
tensorrt_llm/_torch/auto_deploy/models/factory.py	Minor: Removed type annotation from registry method.
tensorrt_llm/_torch/auto_deploy/README.md, examples/auto_deploy/README.md	Major: Expanded and clarified documentation for new modular config and transform pipeline.
tensorrt_llm/_torch/auto_deploy/llm_args.py, .../config/default.yaml	Added default config YAML for modular transforms and updated config classes.
requirements.txt, setup.py	Added omegaconf, updated pydantic-settings to include YAML extra, and ensured YAML files are included in packaging.
Tests (multiple files)	Updated to in-place graph transform semantics, added/expanded quantized MoE, sharding, config, and torch attention backend tests.
examples/auto_deploy/.vscode/launch.json	Updated debug launch args to separate model kwargs and add prompt max tokens; added dry-run option.
examples/auto_deploy/build_and_run_flux.py	Removed reassignment of graph module after quantize and fuse calls, relying on in-place modifications.
tensorrt_llm/_torch/auto_deploy/custom_ops/__init__.py	Added import for torch_backend_attention module.
tensorrt_llm/_torch/auto_deploy/transformations/library/__init__.py	Replaced ep_sharding import with quantize_moe import.
tests/unittest/_torch/auto_deploy/_utils_test/_graph_test_helpers.py	Added option to skip output assertions; added sharding pattern detection test helper.
tests/unittest/_torch/auto_deploy/_utils_test/torch_attention_reference.py	New: Reference implementations for torch backend attention for testing.
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_ad_moe_op.py	Added tests for FP8 and FP4 quantized MoE ops; refactored common setup.
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_attention_op.py	Refactored attention op tests to use centralized torch attention reference.
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_attention_op.py	Refactored flashinfer attention tests to use torch attention reference.
tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_attention_op.py	New: Comprehensive unit tests for torch backend attention op with numpy reference.
tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_moe_fusion.py	Added support and tests for quantized FP8 and FP4 MoE fusion.
tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_quant_moe.py	New: Unit tests for quantize_moe transform with FP8 and NVFP4.
tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_config.py	New: Tests for DynamicYamlMixInForSettings covering YAML loading, merging, precedence, and errors.
tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_bmm_sharding.py	Refactored BMM sharding tests to separate detection and execution; added detection tests.
tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py	Refactored EP sharding tests similarly; added detection tests.
tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py	Refactored TP sharding tests similarly; added detection tests.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant CLI/Script
    participant ConfigManager
    participant InferenceOptimizer
    participant TransformRegistry
    participant Transform
    participant ModelFactory
    participant GraphModule

    User->>CLI/Script: Launch with CLI args/YAML configs
    CLI/Script->>ConfigManager: Load and merge configs (deep merge, YAML, CLI)
    ConfigManager-->>CLI/Script: ExperimentConfig/AutoDeployConfig
    CLI/Script->>ModelFactory: Build model (from config)
    CLI/Script->>InferenceOptimizer: Initialize with ModelFactory and transform config
    InferenceOptimizer->>TransformRegistry: Lookup transforms by stage
    loop For each stage in config
        InferenceOptimizer->>Transform: Instantiate Transform with config
        Transform->>GraphModule: Apply transform (in-place)
        GraphModule-->>Transform: Updated graph and metadata
    end
    InferenceOptimizer-->>CLI/Script: Optimized GraphModule/nn.Module
    CLI/Script->>User: Ready for inference or export

Loading

Suggested reviewers

• qiaoxj07
• juney-nvidia
• pcastonguay

Poem

In the warren of code, a rabbit did find
A garden of transforms, each neatly aligned.
With YAMLs that merge and configs that flex,
The graph now transforms with minimal vex.
Quantized and sharded, so nimble and bright—
This modular pipeline is a bunny's delight!
🐇✨

✨ Finishing Touches

• 📝 Generate Docstrings

---

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Explain this complex logic.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai explain this code block.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.