Skip to content

[Feature]: Refactor dist ops in AutoDeploy #9198

Closed
Task
Closed
[Feature]: Refactor dist ops in AutoDeploy#9198
Task
Assignees
MrGeva
Labels
AutoDeploy<NV> AutoDeploy Backendfeature requestNew feature or request. This includes new model, dtype, functionality support

Description

@lucaslie
lucaslie
opened on Nov 16, 2025

🚀 The feature, motivation and pitch

Right now, AutoDeploy's dist ops handle a lot of dispatch logic:

  1. They check for trtllm dist ops or fall back to torch dist ops
  2. [https://nvbugs/5647400] [fix] Enlarged the AllReduce workspace size to 64MB. Added AllReduce strategy to AD config. #9145 incorporates a torch fallback for a fused all_reduce op
  3. They offer little configurability and instead use preconfigured logic.

Suggestion:

  1. Make dist ops more atomic, e.g., torch.ops.auto_deploy.trtllm_all_reduce should call trt-llm's all-reduce and nothing else. There can be an equivalent torch.ops.auto_deploy.torch_all_reduce that is only calling torch's all_reduce under the hood.
  2. Dispatching can be done on the transform level where the right ops can be inserted

Alternatives

No response

Additional context

No response

Before submitting a new issue...

  • Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

Activity

  1. lucaslie
    assigned
    MrGeva
    on Nov 16, 2025
  2. lucaslie
    added the
    Task
    issue type on Nov 16, 2025
  3. lucaslie
    added
    feature requestNew feature or request. This includes new model, dtype, functionality support
    on Nov 16, 2025
  4. lucaslie
    moved this to Backlog in AutoDeploy Boardon Nov 16, 2025
  5. lucaslie
    added this to AutoDeploy Boardon Nov 16, 2025
  6. lucaslie
    added
    AutoDeploy<NV> AutoDeploy Backend
    on Nov 16, 2025
  7. MrGeva
    moved this from Backlog to In progress in AutoDeploy Boardon Nov 17, 2025

  8. MrGeva commented on Nov 17, 2025

    @MrGeva
    MrGeva
    on Nov 17, 2025
    Collaborator

    Analyzed current code
    Currently these custom ops call distributed ops:

    torch_dist_all_reduce, 
torch_dist_all_gather, 
fused_allreduce_residual_rmsnorm, 
fused_fp8_linear_all_reduce, 
trtllm_dist_fused_linear_all_reduce* (which seems unused).

    They all check is_trtllm_op_available and call the torch/trtllm op accordingly:

    if trtllm_dist.is_trtllm_op_available():
        return trtllm_dist.trtllm_allgather(tensor, dim=dim, sizes=sizes)
    tl = [torch.zeros_like(tensor) for _ in range(dist.get_world_size())]
    dist.all_gather(tl, tensor)

    these ops are used in patterns: _allreduce_residual_rmsnorm_pattern, _allreduce_residual_rmsnorm_pattern2
    and in replace functions: _allreduce_residual_rmsnorm_repl,

    the transforms that insert them are:
    _shard_parameter_node(), _insert_sharded_moe(), _insert_sharded_mxfp4_mlp_ep(). BMMShardingInfo._apply, FuseAllreduceResidualRMSNorm

  9. MrGeva
    moved this from In progress to In review in AutoDeploy Boardon Nov 17, 2025
  10. MrGeva
    added a commit that references this issue on Dec 1, 2025

    [#9198][feat] Refactor dist ops in AutoDeploy (#9301)

    c9771eb
  11. MrGeva
    moved this from In review to Done in AutoDeploy Boardon Dec 2, 2025
  12. MrGeva
    closed this as completedby moving to Done in AutoDeploy Boardon Dec 2, 2025
  13. lucaslie
    added a commit that references this issue on Dec 3, 2025

    [NVIDIA#9198][feat] Refactor dist ops in AutoDeploy (NVIDIA#9301)

    d0c18ae
  14. MinaHuai
    added a commit that references this issue on Dec 10, 2025

    [TRTLLM-7963][feat] Cold L2 cache when doing autotune benchmarking. (N…

    3ff5534
  15. codego7250
    added a commit that references this issue on Dec 11, 2025

    [NVIDIA#9198][feat] Refactor dist ops in AutoDeploy (NVIDIA#9301)

    b92c5d3
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

AutoDeploy<NV> AutoDeploy Backendfeature requestNew feature or request. This includes new model, dtype, functionality support

Type

Task

Projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions