diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml
index 2f26b99b821f..42d4e763d979 100644
--- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml
+++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml
@@ -80,6 +80,7 @@ transforms:
     support_partial_config: true
     sharding_dims: ['tp', 'ep', 'bmm']
     allreduce_strategy: 'AUTO'
+    dist_backend: auto
     requires_shape_prop: true
   sharding_transform_executor:
     stage: sharding
diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md b/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md
index 5a65d6104932..27f18bed1b92 100644
--- a/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md
+++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md
@@ -17,13 +17,12 @@ The table below lists the operators ordered by their backend.
 | `torch.ops.auto_deploy.torch_attention` | Grouped SDPA implementation with `bsnd` and `bnsd` layout supported |
 | `torch.ops.auto_deploy.torch_attention_repeat_kv` | KV repetition for attention |
 | `torch.ops.auto_deploy.torch_attention_sdpa` | Standard SDPA implementation |
-| `torch.ops.auto_deploy.torch_dist_all_gather` | Distributed all-gather operation |
-| `torch.ops.auto_deploy.torch_dist_all_reduce` | Distributed all-reduce operation |
+| `torch.ops.auto_deploy.torch_dist_all_gather` | Distributed all-gather operation (PyTorch backend, demollm mode) |
+| `torch.ops.auto_deploy.torch_dist_all_reduce` | Distributed all-reduce operation (PyTorch backend, demollm mode) |
 | `torch.ops.auto_deploy.torch_linear_simple` | Simple linear layer implementation |
 | `torch.ops.auto_deploy.torch_moe` | Mixture of Experts implementation |
 | `torch.ops.auto_deploy.torch_moe_fused` | Fused Mixture of Experts implementation |
 | `torch.ops.auto_deploy.torch_quant_fn` | Generic quantization function that scales, rounds, and clamps input values |
-| `torch.ops.auto_deploy.torch_quant_fused_fp8_linear_all_reduce` | Fused FP8 linear layer followed by all-reduce operation |
 | `torch.ops.auto_deploy.torch_quant_nvfp4_linear` | FP4 quantized linear layer |
 | `torch.ops.auto_deploy.torch_quant_fp8_linear` | FP8 quantized linear layer |
 | `torch.ops.auto_deploy.torch_rope_with_complex_freqs` | RoPE with complex frequencies |
@@ -38,4 +37,6 @@ The table below lists the operators ordered by their backend.
 | `torch.ops.auto_deploy.triton_rope_on_flattened_inputs` | Triton RoPE on flattened inputs |
 | `torch.ops.auto_deploy.triton_rope_with_input_pos` | Triton RoPE with input positions |
 | `torch.ops.auto_deploy.trtllm_moe_fused` | TensorRT LLM fused MoE implementation |
-| `torch.ops.auto_deploy.trtllm_dist_fused_linear_all_reduce` | TensorRT LLM fused linear layer followed by all-reduce operation |
+| `torch.ops.auto_deploy.trtllm_dist_all_gather` | Distributed all-gather operation (TRT-LLM backend, MPI mode) |
+| `torch.ops.auto_deploy.trtllm_dist_all_reduce` | Distributed all-reduce operation (TRT-LLM backend, MPI mode) |
+| `torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm` | Fused all-reduce + residual add + RMSNorm (TRT-LLM backend, MPI mode) |
diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py
deleted file mode 100644
index 2ff656891d9d..000000000000
--- a/tensorrt_llm/_torch/auto_deploy/custom_ops/dist.py
+++ /dev/null
@@ -1,48 +0,0 @@
-"""Custom ops required for implementing tensor parallelism."""
-
-from typing import List, Optional
-
-import torch
-
-from ..distributed import common as dist
-from ..distributed import trtllm as trtllm_dist
-
-
-@torch.library.custom_op("auto_deploy::torch_dist_all_gather", mutates_args=(), device_types="cuda")
-def all_gather(
-    tensor: torch.Tensor, dim: int = 0, sizes: Optional[List[int]] = None
-) -> torch.Tensor:
-    """All gather followed by concat in dim = 0. This is the default nccl behavior."""
-    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)
-    return torch.cat(tl, dim=dim)
-
-
-@all_gather.register_fake
-def all_gather_fake(tensor, dim=0):
-    return torch.cat([torch.empty_like(tensor) for _ in range(dist.get_world_size())], dim=dim)
-
-
-@torch.library.custom_op("auto_deploy::torch_dist_all_reduce", mutates_args=(), device_types="cuda")
-def all_reduce(t: torch.Tensor, strategy: str) -> torch.Tensor:
-    """All_reduce across the ranks. Reduction op is SUM. Strategy is MANDATORY.
-
-    Args:
-        t: Tensor to reduce across ranks
-        strategy: AllReduce strategy - "AUTO", "NCCL", "ONESHOT", "TWOSHOT", "MIN_LATENCY", etc.
-
-    NOTE: this op requires an extra memory copy and should ONLY be used for debugging + testing. For
-    efficient all_reduce ops one should write/replace it with a fused op.
-    """
-    if trtllm_dist.is_trtllm_op_available():
-        return trtllm_dist.trtllm_allreduce(t, op=dist.ReduceOp.SUM, strategy=strategy)
-    t_res = t.clone()
-    dist.all_reduce(t_res, op=dist.ReduceOp.SUM)
-    return t_res
-
-
-@all_reduce.register_fake
-def all_reduce_fake(tensor, strategy):
-    return torch.empty_like(tensor)
diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear.py
index 8ecd42150175..214626ad24e1 100644
--- a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear.py
+++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear.py
@@ -4,9 +4,6 @@
 
 import torch
 
-from ..distributed import common as dist
-from ..distributed import trtllm as trtllm_dist
-
 
 @torch.library.custom_op("auto_deploy::torch_linear_simple", mutates_args=())
 def simple(input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor]) -> torch.Tensor:
@@ -24,26 +21,4 @@ def simple(input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tenso
 @simple.register_fake
 def simple_fake(input, weight, bias):
     """Fake implementation of simple_linear."""
-    # return torch.empty(
-    # input.shape[:-1] + (weight.shape[-1],), dtype=input.dtype, device=input.device
-    # )
-    return torch.ops.aten.linear(input, weight, bias)
-
-
-@torch.library.custom_op(
-    "auto_deploy::trtllm_dist_fused_linear_all_reduce", mutates_args=(), device_types="cuda"
-)
-def fused_linear_all_reduce(
-    input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], strategy: str
-) -> torch.Tensor:
-    """Fused linear followed by all_reduce on the output. Strategy is MANDATORY."""
-    output = torch.ops.aten.linear(input, weight, bias)
-    if trtllm_dist.is_trtllm_op_available():
-        return trtllm_dist.trtllm_allreduce(output, op=dist.ReduceOp.SUM, strategy=strategy)
-    dist.all_reduce(output, op=dist.ReduceOp.SUM)
-    return output
-
-
-@fused_linear_all_reduce.register_fake
-def fused_linear_all_reduce_fake(input, weight, bias, strategy):
     return torch.ops.aten.linear(input, weight, bias)
diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py
index 7c1dbc6f0ae1..4d9c7949644d 100644
--- a/tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py
+++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py
@@ -7,8 +7,6 @@
 from flashinfer import bmm_fp8
 from torch import nn
 
-from ..distributed import common as dist
-from ..distributed import trtllm as trtllm_dist
 from .torch_libs.float8_python_api import addmm_float8_unwrapped
 
 TRTLLM_FP4_OP_AVAILABLE = True
@@ -238,39 +236,6 @@ def fp8_linear_fake(
     return torch.ops.aten.linear(input, weight_fp8.to(input.dtype), bias)
 
 
-@torch.library.custom_op("auto_deploy::torch_quant_fused_fp8_linear_all_reduce", mutates_args=())
-@torch.compile(dynamic=True)
-def fused_fp8_linear_all_reduce(
-    input: torch.Tensor,
-    weight_fp8: torch.Tensor,
-    strategy: str,
-    bias: Optional[torch.Tensor] = None,
-    input_scale: Optional[torch.Tensor] = None,
-    weight_scale: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    out = torch.ops.auto_deploy.torch_quant_fp8_linear(
-        input, weight_fp8, bias, input_scale, weight_scale
-    )
-    if trtllm_dist.is_trtllm_op_available():
-        return trtllm_dist.trtllm_allreduce(out, op=dist.ReduceOp.SUM, strategy=strategy)
-    dist.all_reduce(out, op=dist.ReduceOp.SUM)
-    return out
-
-
-@fused_fp8_linear_all_reduce.register_fake
-def fused_fp8_linear_all_reduce_fake(
-    input: torch.Tensor,
-    weight_fp8: torch.Tensor,
-    strategy: str,
-    bias: Optional[torch.Tensor] = None,
-    input_scale: Optional[torch.Tensor] = None,
-    weight_scale: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    return torch.ops.auto_deploy.torch_quant_fp8_linear(
-        input, weight_fp8, bias, input_scale, weight_scale
-    )
-
-
 class FP8Linear(nn.Linear):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_dist.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_dist.py
new file mode 100644
index 000000000000..5cbda619704a
--- /dev/null
+++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/torch_dist.py
@@ -0,0 +1,52 @@
+"""Custom ops required for implementing tensor parallelism.
+
+This module defines atomic distributed ops - each op uses a specific backend
+(torch.distributed or TRT-LLM) without internal dispatch logic.
+"""
+
+from typing import List, Optional
+
+import torch
+
+from ..distributed import common as dist
+
+# ============================================================================
+# PyTorch Distributed Backend Ops (demollm mode)
+# ============================================================================
+
+
+@torch.library.custom_op("auto_deploy::torch_dist_all_gather", mutates_args=(), device_types="cuda")
+def torch_dist_all_gather(
+    tensor: torch.Tensor, dim: int = 0, sizes: Optional[List[int]] = None
+) -> torch.Tensor:
+    """All gather using PyTorch distributed backend.
+
+    This op always uses torch.distributed.all_gather and is used in demollm mode.
+    """
+    tl = [torch.zeros_like(tensor) for _ in range(dist.get_world_size())]
+    dist.all_gather(tl, tensor)
+    return torch.cat(tl, dim=dim)
+
+
+@torch_dist_all_gather.register_fake
+def torch_dist_all_gather_fake(tensor, dim=0, sizes=None):
+    return torch.cat([torch.empty_like(tensor) for _ in range(dist.get_world_size())], dim=dim)
+
+
+@torch.library.custom_op("auto_deploy::torch_dist_all_reduce", mutates_args=(), device_types="cuda")
+def torch_dist_all_reduce(t: torch.Tensor, strategy: str) -> torch.Tensor:
+    """All_reduce using PyTorch distributed backend. Reduction op is SUM.
+
+    This op always uses torch.distributed.all_reduce and is used in demollm mode.
+
+    NOTE: this op requires an extra memory copy and should ONLY be used for debugging + testing. For
+    efficient all_reduce ops one should write/replace it with a fused op.
+    """
+    t_res = t.clone()
+    dist.all_reduce(t_res, op=dist.ReduceOp.SUM)
+    return t_res
+
+
+@torch_dist_all_reduce.register_fake
+def torch_dist_all_reduce_fake(tensor, strategy):
+    return torch.empty_like(tensor)
diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/trtllm_dist.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/trtllm_dist.py
new file mode 100644
index 000000000000..1f4b9229ca4e
--- /dev/null
+++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/trtllm_dist.py
@@ -0,0 +1,134 @@
+"""TRT-LLM distributed operations and fused kernels.
+
+This module defines atomic TRT-LLM-specific ops that use optimized kernels.
+The torch fallback variants are defined separately to enable multi-pattern matching.
+"""
+
+from typing import List, Optional
+
+import torch
+
+# use trtllm distributed ops to improve TP performance if possible
+from ....mapping import Mapping
+from ...distributed import AllReduce, allgather
+from ...modules.linear import AllReduceFusionOp, AllReduceParams, AllReduceStrategy
+from ..distributed.common import ReduceOp, get_rank_world_size, get_world_size, is_ompi
+
+# Cache AllReduce modules to avoid recreating on every call
+# This is critical for CUDA graph compatibility - recreating modules during
+# warmup causes hangs due to workspace allocation with CPU synchronization
+_allreduce_cache = {}
+
+
+def trtllm_allgather(tensor, dim, sizes=None):
+    rank, world_size = get_rank_world_size()
+    p_config = Mapping(world_size=world_size, tp_size=world_size, rank=rank)
+    return allgather(tensor, p_config, dim=dim, sizes=sizes)
+
+
+def trtllm_allreduce(tensor, op, strategy: str, all_reduce_params=None):
+    rank, world_size = get_rank_world_size()
+    assert op == ReduceOp.SUM, "TRT-LLM all reduce only supports SUM op."
+
+    # Convert string strategy to enum
+    try:
+        strategy_enum = getattr(AllReduceStrategy, strategy)
+    except AttributeError:
+        raise ValueError(
+            f"Invalid allreduce strategy: {strategy}. "
+            f"Valid options: AUTO, NCCL, ONESHOT, TWOSHOT, MIN_LATENCY, "
+            f"LOWPRECISION, UB, MNNVL, NCCL_SYMMETRIC"
+        )
+
+    # Cache key includes rank, world_size, dtype, and strategy to handle different configurations
+    cache_key = (rank, world_size, tensor.dtype, strategy_enum)
+    if cache_key not in _allreduce_cache:
+        p_config = Mapping(world_size=world_size, tp_size=world_size, rank=rank)
+        _allreduce_cache[cache_key] = AllReduce(
+            mapping=p_config, strategy=strategy_enum, dtype=tensor.dtype
+        )
+
+    torch_op = _allreduce_cache[cache_key]
+    return torch_op(tensor, all_reduce_params=all_reduce_params)
+
+
+# ============================================================================
+# TRT-LLM Backend Ops (MPI mode)
+# ============================================================================
+
+
+@torch.library.custom_op(
+    "auto_deploy::trtllm_dist_all_gather", mutates_args=(), device_types="cuda"
+)
+def trtllm_dist_all_gather(
+    tensor: torch.Tensor, dim: int = 0, sizes: Optional[List[int]] = None
+) -> torch.Tensor:
+    """All gather using TRT-LLM optimized backend.
+
+    This op always uses TRT-LLM's optimized allgather and is used in MPI mode.
+    """
+    return trtllm_allgather(tensor, dim=dim, sizes=sizes)
+
+
+@trtllm_dist_all_gather.register_fake
+def trtllm_dist_all_gather_fake(tensor, dim=0, sizes=None):
+    return torch.cat([torch.empty_like(tensor) for _ in range(get_world_size())], dim=dim)
+
+
+@torch.library.custom_op(
+    "auto_deploy::trtllm_dist_all_reduce", mutates_args=(), device_types="cuda"
+)
+def trtllm_dist_all_reduce(t: torch.Tensor, strategy: str) -> torch.Tensor:
+    """All_reduce using TRT-LLM optimized backend. Reduction op is SUM.
+
+    This op always uses TRT-LLM's optimized allreduce and is used in MPI mode.
+    """
+    return trtllm_allreduce(t, op=ReduceOp.SUM, strategy=strategy)
+
+
+@trtllm_dist_all_reduce.register_fake
+def trtllm_dist_all_reduce_fake(tensor, strategy):
+    return torch.empty_like(tensor)
+
+
+# TRT-LLM fused op (atomic - always uses TRT-LLM backend)
+@torch.library.custom_op(
+    "dist::trtllm_fused_allreduce_residual_rmsnorm", mutates_args=(), device_types="cuda"
+)
+def trtllm_fused_allreduce_residual_rmsnorm(
+    tensor: torch.Tensor,
+    residual: torch.Tensor,
+    norm_weight: torch.Tensor,
+    eps: float,
+    strategy: str,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Fused allreduce + residual + rmsnorm using TRT-LLM optimized kernel.
+
+    This op always uses TRT-LLM's fused kernel and is used in MPI mode.
+    """
+    all_reduce_params = AllReduceParams(
+        fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
+        bias=None,
+        residual=residual,
+        norm_weight=norm_weight,
+        eps=eps,
+    )
+    return trtllm_allreduce(
+        tensor, ReduceOp.SUM, strategy=strategy, all_reduce_params=all_reduce_params
+    )
+
+
+@trtllm_fused_allreduce_residual_rmsnorm.register_fake
+def trtllm_fused_allreduce_residual_rmsnorm_fake(
+    tensor: torch.Tensor,
+    residual: torch.Tensor,
+    norm_weight: torch.Tensor,
+    eps: float,
+    strategy: str,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return torch.empty_like(tensor), torch.empty_like(tensor)
+
+
+def is_trtllm_op_available():
+    """Check if TRT-LLM ops are available and running with MPI."""
+    return is_ompi()
diff --git a/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py b/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py
deleted file mode 100644
index 386083f7cc6e..000000000000
--- a/tensorrt_llm/_torch/auto_deploy/distributed/trtllm.py
+++ /dev/null
@@ -1,119 +0,0 @@
-import torch
-
-from .common import ReduceOp, get_rank_world_size, is_ompi
-
-# use trtllm distributed ops to improve TP performance if possible
-try:
-    from ....mapping import Mapping
-    from ...distributed import AllReduce, allgather
-    from ...modules.linear import AllReduceFusionOp, AllReduceParams, AllReduceStrategy
-
-    # Cache AllReduce modules to avoid recreating on every call
-    # This is critical for CUDA graph compatibility - recreating modules during
-    # warmup causes hangs due to workspace allocation with CPU synchronization
-    _allreduce_cache = {}
-
-    def trtllm_allgather(tensor, dim, sizes=None):
-        rank, world_size = get_rank_world_size()
-        p_config = Mapping(world_size=world_size, tp_size=world_size, rank=rank)
-        return allgather(tensor, p_config, dim=dim, sizes=sizes)
-
-    def trtllm_allreduce(tensor, op, strategy: str, all_reduce_params=None):
-        rank, world_size = get_rank_world_size()
-        assert op == ReduceOp.SUM, "TRT-LLM all reduce only supports SUM op."
-
-        # Convert string strategy to enum
-        try:
-            strategy_enum = getattr(AllReduceStrategy, strategy)
-        except AttributeError:
-            raise ValueError(
-                f"Invalid allreduce strategy: {strategy}. "
-                f"Valid options: AUTO, NCCL, ONESHOT, TWOSHOT, MIN_LATENCY, "
-                f"LOWPRECISION, UB, MNNVL, NCCL_SYMMETRIC"
-            )
-
-        # Cache key includes rank, world_size, dtype, and strategy to handle different configurations
-        cache_key = (rank, world_size, tensor.dtype, strategy_enum)
-        if cache_key not in _allreduce_cache:
-            p_config = Mapping(world_size=world_size, tp_size=world_size, rank=rank)
-            _allreduce_cache[cache_key] = AllReduce(
-                mapping=p_config, strategy=strategy_enum, dtype=tensor.dtype
-            )
-
-        torch_op = _allreduce_cache[cache_key]
-        return torch_op(tensor, all_reduce_params=all_reduce_params)
-
-    @torch.library.custom_op(
-        "dist::fused_allreduce_residual_rmsnorm", mutates_args=(), device_types="cuda"
-    )
-    def fused_allreduce_residual_rmsnorm(
-        tensor: torch.Tensor,
-        residual: torch.Tensor,
-        norm_weight: torch.Tensor,
-        eps: float,
-        strategy: str = "AUTO",
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Fusing allreduce, residual (add), and hf_rms_norm together.
-
-        When TRT-LLM ops are available (MPI mode), uses the fused kernel.
-        Otherwise, falls back to separate operations using torch distributed.
-        """
-        # Only use TRT-LLM fused op when running with MPI
-        if is_trtllm_op_available():
-            all_reduce_params = AllReduceParams(
-                fusion_op=AllReduceFusionOp.RESIDUAL_RMS_NORM,
-                bias=None,
-                residual=residual,
-                norm_weight=norm_weight,
-                eps=eps,
-            )
-            return trtllm_allreduce(
-                tensor, ReduceOp.SUM, strategy=strategy, all_reduce_params=all_reduce_params
-            )
-        else:
-            # Fallback: unfused implementation using torch distributed
-            # This is used in demollm mode without MPI
-            from .common import all_reduce as torch_all_reduce
-
-            # 1. All-reduce the tensor
-            tensor_reduced = tensor.clone()
-            torch_all_reduce(tensor_reduced, op=ReduceOp.SUM)
-
-            # 2. Add residual
-            tensor_with_residual = tensor_reduced + residual
-
-            # 3. Apply RMSNorm using PyTorch's built-in function
-            norm_out = torch.nn.functional.rms_norm(
-                tensor_with_residual,
-                normalized_shape=(tensor_with_residual.size(-1),),
-                weight=norm_weight,
-                eps=eps,
-            )
-
-            return norm_out, tensor_with_residual
-
-    @fused_allreduce_residual_rmsnorm.register_fake
-    def fused_allreduce_residual_rmsnorm_fake(
-        tensor: torch.Tensor,
-        residual: torch.Tensor,
-        norm_weight: torch.Tensor,
-        eps: float,
-        strategy: str = "AUTO",
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        return torch.empty_like(tensor), torch.empty_like(tensor)
-
-    TRTLLM_OP_AVAILABLE = True
-except ImportError:
-
-    def trtllm_allgather(tensor, dim, sizes=None):
-        raise ImportError("TRT-LLM is not available.")
-
-    def trtllm_allreduce(tensor, op, strategy: str, all_reduce_params=None):
-        raise ImportError("TRT-LLM is not available.")
-
-    TRTLLM_OP_AVAILABLE = False
-
-
-def is_trtllm_op_available():
-    # TRT-LLM only work with MPI
-    return TRTLLM_OP_AVAILABLE and is_ompi()
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/collectives.py b/tensorrt_llm/_torch/auto_deploy/transform/library/collectives.py
index 97d4e4bbbeb0..6eb5371f4026 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/collectives.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/collectives.py
@@ -1,3 +1,11 @@
+"""Transformations for fusing collective operations.
+
+This module registers TRT-LLM backend patterns only. Fusion is only applied
+when TRT-LLM is available (MPI mode) since it provides optimized fused kernels.
+The torch backend (demollm mode) does not benefit from fusion.
+"""
+
+from functools import partial
 from typing import Tuple
 
 import torch
@@ -16,63 +24,77 @@
 # * ...
 
 
-def _allreduce_residual_rmsnorm_pattern(
-    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 0.1253
-):
-    """
-    Reference PyTorch composition of:
-        y = all_reduce(x)
-        z = residual + y
-        normed = RMSNorm(z, weight, eps)
-    Returns (normed, z)
-    """
+# ============================================================================
+# Pattern Template Factory Functions
+# ============================================================================
 
-    input_dtype = x.dtype
-    hidden_states = torch.ops.auto_deploy.torch_dist_all_reduce(x, "AUTO")
-    add = residual + hidden_states
 
-    hidden_states = add.to(torch.float32)
-    variance = hidden_states.pow(2).mean(-1, keepdim=True)
-    hidden_states = hidden_states * torch.rsqrt(variance + eps)
+def _make_allreduce_residual_rmsnorm_pattern(
+    add_order: str = "residual_first", strategy: str = "AUTO"
+):
+    """Factory function to create pattern functions for allreduce+residual+rmsnorm fusion.
 
-    normed = weight * hidden_states.to(input_dtype)
+    Args:
+        add_order: Either "residual_first" (residual + x) or "x_first" (x + residual)
+        strategy: AllReduce strategy to use in the pattern
 
-    return normed, add
+    Returns:
+        A pattern function that can be used with register_ad_pattern
+    """
 
+    def pattern_fn(
+        x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 0.1253
+    ):
+        """Pattern: trtllm_dist_all_reduce(x) -> add residual -> RMSNorm
 
-def _allreduce_residual_rmsnorm_pattern2(
-    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 0.1253
-):
-    """
-    Reference PyTorch composition of:
-        y = all_reduce(x)
-        z = y + residual
-        normed = RMSNorm(z, weight, eps)
-    Returns (normed, z)
-    """
+        Reference PyTorch composition:
+            y = trtllm_dist_all_reduce(x)
+            z = residual + y  (or y + residual)
+            normed = RMSNorm(z, weight, eps)
+        Returns (normed, z)
+        """
+        input_dtype = x.dtype
+        hidden_states = torch.ops.auto_deploy.trtllm_dist_all_reduce(x, strategy)
+
+        # Handle addition order
+        if add_order == "residual_first":
+            add = residual + hidden_states
+        else:  # x_first
+            add = hidden_states + residual
 
-    input_dtype = x.dtype
-    hidden_states = torch.ops.auto_deploy.torch_dist_all_reduce(x, "AUTO")
-    add = hidden_states + residual
+        hidden_states = add.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + eps)
 
-    hidden_states = add.to(torch.float32)
-    variance = hidden_states.pow(2).mean(-1, keepdim=True)
-    hidden_states = hidden_states * torch.rsqrt(variance + eps)
+        normed = weight * hidden_states.to(input_dtype)
 
-    normed = weight * hidden_states.to(input_dtype)
+        return normed, add
 
-    return normed, add
+    return pattern_fn
 
 
-def _allreduce_residual_rmsnorm_repl(
-    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float
+def _allreduce_residual_rmsnorm_replacement(
+    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float, strategy: str
 ):
-    return torch.ops.dist.fused_allreduce_residual_rmsnorm(x, residual, weight, eps)
+    """Replacement using TRT-LLM fused kernel."""
+    return torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm(
+        x, residual, weight, eps, strategy
+    )
+
+
+# ============================================================================
+# Transform Implementation
+# ============================================================================
 
 
 @TransformRegistry.register("fuse_allreduce_residual_rmsnorm")
 class FuseAllreduceResidualRMSNorm(BaseTransform):
-    """Fuse (allreduce + residual add + RMSNorm) into one fused op with tuple output."""
+    """Fuse (allreduce + residual add + RMSNorm) into one fused op with tuple output.
+
+    This transform only applies when TRT-LLM ops are used (MPI mode), as it provides
+    optimized fused kernels. The torch backend (demollm mode) does not benefit from
+    this fusion and uses unfused operations.
+    """
 
     def _apply(
         self,
@@ -92,21 +114,43 @@ def _apply(
             0.1253,  # eps
         ]
 
+        op_ignore_types = {torch.ops.aten.to.dtype: (torch.dtype,)}
+        scalar_workaround = {"eps": 0.1253}
+
+        # ============================================================================
+        # Instantiate Pattern Functions
+        # ============================================================================
+
+        # Get the allreduce strategy from shared_config
+        strategy = shared_config.sharding_transform_container.allreduce_strategy.name
+
+        # TRT-LLM backend (MPI mode) - two patterns for different addition orders
+        _allreduce_residual_rmsnorm_pattern_trtllm = _make_allreduce_residual_rmsnorm_pattern(
+            add_order="residual_first", strategy=strategy
+        )
+        _allreduce_residual_rmsnorm_pattern2_trtllm = _make_allreduce_residual_rmsnorm_pattern(
+            add_order="x_first", strategy=strategy
+        )
+
+        # Register TRT-LLM backend patterns only (no torch backend fusion)
+        # Pattern 1: residual + allreduce(x)
         register_ad_pattern(
-            search_fn=_allreduce_residual_rmsnorm_pattern,
-            replace_fn=_allreduce_residual_rmsnorm_repl,
+            search_fn=_allreduce_residual_rmsnorm_pattern_trtllm,
+            replace_fn=partial(_allreduce_residual_rmsnorm_replacement, strategy=strategy),
             patterns=patterns,
             dummy_args=dummy_args,
-            op_ignore_types={torch.ops.aten.to.dtype: (torch.dtype,)},
-            scalar_workaround={"eps": 0.1253},
+            op_ignore_types=op_ignore_types,
+            scalar_workaround=scalar_workaround,
         )
+
+        # Pattern 2: allreduce(x) + residual
         register_ad_pattern(
-            search_fn=_allreduce_residual_rmsnorm_pattern2,
-            replace_fn=_allreduce_residual_rmsnorm_repl,
+            search_fn=_allreduce_residual_rmsnorm_pattern2_trtllm,
+            replace_fn=partial(_allreduce_residual_rmsnorm_replacement, strategy=strategy),
             patterns=patterns,
             dummy_args=dummy_args,
-            op_ignore_types={torch.ops.aten.to.dtype: (torch.dtype,)},
-            scalar_workaround={"eps": 0.1253},
+            op_ignore_types=op_ignore_types,
+            scalar_workaround=scalar_workaround,
         )
 
         num_matches = patterns.apply(gm.graph)
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
index 1ed5652c6287..6230cd3031d9 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
@@ -40,6 +40,7 @@
 )
 from ...utils.sharding_utils import (
     BMMShardingInfo,
+    DistBackend,
     EPShardingInfo,
     LayerType,
     ParameterUpdateInfo,
@@ -92,6 +93,8 @@ def _validate_allreduce_strategy(cls, v):
         """Convert string names like 'AUTO' to AllReduceStrategy enum."""
         return validate_allreduce_strategy(v)
 
+    dist_backend: DistBackend = Field(default=DistBackend.AUTO)
+
 
 @TransformRegistry.register("sharding_transform_executor")
 class ShardingTransformExecutor(BaseTransform):
@@ -1004,7 +1007,8 @@ def detect_dp_bmm_shard(
         base_size = bmm_batch_size // world_size
         remainder = bmm_batch_size % world_size
 
-        # NOTE: our torch.ops.auto_deploy.torch_dist_all_gather doesn't support uneven splits at the moment.
+        # NOTE: our torch.ops.auto_deploy.torch_dist_all_gather/trtllm_dist_all_gather
+        #  doesn't support uneven splits at the moment.
         if remainder:
             ad_logger.warning(
                 f"BMM batch size {bmm_batch_size} is not divisible by world size {world_size}. "
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/visualization.py b/tensorrt_llm/_torch/auto_deploy/transform/library/visualization.py
index 0ed97c3ad24d..1527ed24149b 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/visualization.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/visualization.py
@@ -77,9 +77,9 @@ def add_outputs_metadata(self, fx_node: torch.fx.node.Node, node: GraphNode):
 # TODO(yudong): make custom_ops configurable
 CUSTOM_OPS = (
     torch.ops.auto_deploy.torch_dist_all_reduce.default,
+    torch.ops.auto_deploy.trtllm_dist_all_reduce.default,
     torch.ops.aten.slice.Tensor,
     torch.ops.auto_deploy.triton_attention_fused_mha_with_cache.default,
-    torch.ops.auto_deploy.trtllm_dist_fused_linear_all_reduce.default,
     torch.ops.auto_deploy.torch_linear_simple.default,
     torch.ops.aten.split_with_sizes.default,
 )
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
index 33c31a9ba223..9cdb20422fb5 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
+++ b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
@@ -305,10 +305,14 @@ def is_bmm_op(node: Node) -> bool:
 
 
 def is_dist_op(node: Node) -> bool:
-    """Check if the node is a distributed op."""
+    """Check if the node is a distributed op (torch or trtllm backend)."""
     dist_ops = {
+        # PyTorch backend ops
         torch.ops.auto_deploy.torch_dist_all_gather,
         torch.ops.auto_deploy.torch_dist_all_reduce,
+        # TRT-LLM backend ops
+        torch.ops.auto_deploy.trtllm_dist_all_gather,
+        torch.ops.auto_deploy.trtllm_dist_all_reduce,
     }
     return is_op(node, dist_ops)
 
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py
index 08441632c396..d12b99545440 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py
+++ b/tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py
@@ -63,6 +63,49 @@ def validate_allreduce_strategy(v):
     return v  # Let Pydantic handle other types
 
 
+def _get_dist_ops(backend: str):
+    """Get the appropriate distributed ops based on backend availability.
+
+    Args:
+        backend: The distributed backend to use. Can be 'auto', 'trtllm', or 'torch'.
+                 'auto' will automatically select based on availability.
+
+    Returns tuple of (all_gather_op, all_reduce_op) for the current backend.
+    """
+    from ..custom_ops.trtllm_dist import is_trtllm_op_available
+
+    # Handle DistBackend enum or string
+    if hasattr(backend, "value"):
+        backend = backend.value
+
+    if backend == "trtllm":
+        # Force TRT-LLM ops
+        return (
+            torch.ops.auto_deploy.trtllm_dist_all_gather.default,
+            torch.ops.auto_deploy.trtllm_dist_all_reduce.default,
+        )
+    elif backend == "torch":
+        # Force PyTorch distributed ops
+        return (
+            torch.ops.auto_deploy.torch_dist_all_gather.default,
+            torch.ops.auto_deploy.torch_dist_all_reduce.default,
+        )
+    else:  # auto
+        # Automatically select based on availability
+        if is_trtllm_op_available():
+            # Use TRT-LLM optimized ops in MPI mode
+            return (
+                torch.ops.auto_deploy.trtllm_dist_all_gather.default,
+                torch.ops.auto_deploy.trtllm_dist_all_reduce.default,
+            )
+        else:
+            # Use PyTorch distributed ops in demollm mode
+            return (
+                torch.ops.auto_deploy.torch_dist_all_gather.default,
+                torch.ops.auto_deploy.torch_dist_all_reduce.default,
+            )
+
+
 def _load_hook(
     state_dict,
     prefix,
@@ -265,6 +308,7 @@ def _insert_sharded_mamba(
     rank: int,
     world_size: int,
     allreduce_strategy: AllReduceStrategy,
+    dist_backend: str,
     add_dist: bool = False,
     min_local_shape: int = 1,
     weights_to_shard: Optional[list[str]] = None,
@@ -379,6 +423,7 @@ def _insert_sharded_mamba(
         dim=SplitDimension.COLUMN,
         rank=rank,
         world_size=world_size,
+        dist_backend=dist_backend,
         add_dist=False,
         min_local_shape=min_local_shape,
         fused_weight_dims=entry_fused_dims,
@@ -444,6 +489,7 @@ def _shard_parameter_node(
     rank: int,
     world_size: int,
     allreduce_strategy: AllReduceStrategy,
+    dist_backend: str,
     add_dist: bool = False,
     min_local_shape: int = 1,
     fused_weight_dims: Optional[list] = None,
@@ -534,19 +580,17 @@ def _shard_parameter_node(
             )
         return
 
-    # figure out the right dist op
-    if dim == 0:
-        # Column split -> all_gather
-        fn_dist = torch.ops.auto_deploy.torch_dist_all_gather.default
-        dist_args = (node, -1)
-    else:
-        # Row split -> all_reduce with strategy
-        fn_dist = torch.ops.auto_deploy.torch_dist_all_reduce.default
-        dist_args = (node, allreduce_strategy.name)
+    # figure out the right dist op (backend-aware)
+    all_gather_op, all_reduce_op = _get_dist_ops(dist_backend)
+    dist_lookup = {
+        0: (all_gather_op, -1),
+        1: (all_reduce_op, allreduce_strategy.name),
+    }
+    fn_dist, *dist_args = dist_lookup[dim]
 
     # add reduction node
     with gm.graph.inserting_after(node):
-        dist_node = gm.graph.call_function(fn_dist, args=dist_args)
+        dist_node = gm.graph.call_function(fn_dist, args=(node,) + tuple(dist_args))
         node.replace_all_uses_with(dist_node)
         dist_node.replace_input_with(dist_node, node)
 
@@ -630,6 +674,7 @@ class WeightShardingInfo(ShardingTransformInfo):
     # used for TP sharding of fused weights
     fused_weight_dims: Optional[list] = None
     allreduce_strategy: Optional[AllReduceStrategy] = None  # Set by ShardingConfig.add() if None
+    dist_backend: Optional[str] = None  # Set by ShardingConfig.add() if None
 
     def quantization_cb(
         self,
@@ -679,6 +724,7 @@ def apply(self, gm: GraphModule, node: Node) -> None:
                 dim=self.split_dim.value,
                 rank=self.rank,
                 world_size=self.world_size,
+                dist_backend=self.dist_backend,
                 add_dist=self.dist_op is not None,
                 min_local_shape=self.min_local_shape,
                 fused_weight_dims=self.fused_weight_dims
@@ -694,6 +740,7 @@ def apply(self, gm: GraphModule, node: Node) -> None:
                 dim=self.split_dim.value,
                 rank=self.rank,
                 world_size=self.world_size,
+                dist_backend=self.dist_backend,
                 add_dist=self.dist_op is not None,
                 min_local_shape=self.min_local_shape,
                 fused_weight_dims=self.fused_weight_dims,
@@ -894,6 +941,7 @@ class BMMShardingInfo(ShardingTransformInfo):
     world_size: int
     start_idx: int
     end_idx: int
+    dist_backend: Optional[str] = None  # Set by ShardingConfig.add() if None
 
     def validate(self, gm: GraphModule = None, node: Node = None) -> bool:
         """Validate the transformation configuration."""
@@ -915,7 +963,8 @@ def validate(self, gm: GraphModule = None, node: Node = None) -> bool:
         # Check if the distribution is balanced
         remainder = bmm_batch_size % self.world_size
 
-        # NOTE: our torch.ops.auto_deploy.torch_dist_all_gather doesn't support uneven splits at the moment.
+        # NOTE: our torch.ops.auto_deploy.torch_dist_all_gather/trtllm_dist_all_gather
+        #  doesn't support uneven splits at the moment.
         if remainder:
             ad_logger.warning(
                 f"BMM batch size {bmm_batch_size} is not divisible by world size {self.world_size}. "
@@ -980,9 +1029,10 @@ def slice_tensor(t: torch.Tensor) -> torch.Tensor:
         handle_tensor(node, rhs_tensor, 1, self.start_idx, self.end_idx)
 
         # Add all_gather node after BMM to collect results
+        all_gather_op, _ = _get_dist_ops(self.dist_backend)
         with gm.graph.inserting_after(node):
             gather_node = gm.graph.call_function(
-                torch.ops.auto_deploy.torch_dist_all_gather.default,
+                all_gather_op,
                 args=(node, 0),  # Gather along batch dimension (0)
             )
             node.replace_all_uses_with(gather_node)
@@ -995,6 +1045,7 @@ def _insert_sharded_moe(
     rank: int,
     world_size: int,
     allreduce_strategy: AllReduceStrategy,
+    dist_backend: str,
     scale_names: Sequence[str] = (),
 ):
     """Update the torch_moe node with sharded weight lists,
@@ -1073,11 +1124,9 @@ def get_partition(lst, world_size, rank):
     node.args = tuple(args)
 
     # -- add an all_reduce node --
+    _, all_reduce_op = _get_dist_ops(dist_backend)
     with gm.graph.inserting_after(node):
-        dist_node = gm.graph.call_function(
-            torch.ops.auto_deploy.torch_dist_all_reduce.default,
-            args=(node, allreduce_strategy.name),
-        )
+        dist_node = gm.graph.call_function(all_reduce_op, args=(node, allreduce_strategy.name))
         node.replace_all_uses_with(dist_node)
         dist_node.replace_input_with(dist_node, node)
 
@@ -1106,11 +1155,12 @@ def _insert_sharded_mxfp4_mlp_ep(
     rank: int,
     world_size: int,
     allreduce_strategy: AllReduceStrategy,
+    dist_backend: str,
 ):
     """Transform a call to auto_deploy::triton_mxfp4_moe into:
       - sharded expert parameters along dim 0 (this rank slice),
       - call to auto_deploy::triton_mxfp4_moe_ep(..., local_lo, local_hi),
-      - followed by torch_dist_all_reduce.
+      - followed by torch_dist_all_reduce/trtllm_dist_all_reduce.
 
     NOTE: allreduce_strategy is MANDATORY and must be explicitly provided.
 
@@ -1152,10 +1202,9 @@ def _insert_sharded_mxfp4_mlp_ep(
     node.args = args_ep
 
     # Add a dist all-reduce after the op (sum partial results across EP ranks)
+    _, all_reduce_op = _get_dist_ops(dist_backend)
     with gm.graph.inserting_after(node):
-        red = gm.graph.call_function(
-            torch.ops.auto_deploy.torch_dist_all_reduce, args=(node, allreduce_strategy.name)
-        )
+        red = gm.graph.call_function(all_reduce_op, args=(node, allreduce_strategy.name))
         node.replace_all_uses_with(red)
         # keep dataflow: red(input=node)
         red.replace_input_with(red, node)
@@ -1164,11 +1213,13 @@ def _insert_sharded_mxfp4_mlp_ep(
 class EPShardingInfo(ShardingTransformInfo):
     """Configuration for EP sharding transformations.
 
-    NOTE: allreduce_strategy will be automatically injected by ShardingConfig.add()
-    if not provided at creation time. The strategy comes from the parent ShardingConfig.
+    NOTE: allreduce_strategy and dist_backend will be automatically injected by
+    ShardingConfig.add() if not provided at creation time. The values come from
+    the parent ShardingConfig.
     """
 
     allreduce_strategy: Optional[AllReduceStrategy] = None  # Set by ShardingConfig.add() if None
+    dist_backend: Optional[str] = None  # Set by ShardingConfig.add() if None
 
     @classmethod
     def from_node(cls, node: Node, **kwargs) -> "EPShardingInfo":
@@ -1187,7 +1238,9 @@ def validate(self, gm: GraphModule = None, node: Node = None) -> bool:
 
     def apply(self, gm: GraphModule, node: Node) -> None:
         """Apply EP sharding transformation to the graph module."""
-        _insert_sharded_moe(gm, node, self.rank, self.world_size, self.allreduce_strategy, [])
+        _insert_sharded_moe(
+            gm, node, self.rank, self.world_size, self.allreduce_strategy, self.dist_backend, []
+        )
 
 
 class MXFP4EPShardingInfo(EPShardingInfo):
@@ -1201,7 +1254,9 @@ def validate(self, gm: GraphModule = None, node: Node = None) -> bool:
         return True
 
     def apply(self, gm: GraphModule, node: Node) -> None:
-        _insert_sharded_mxfp4_mlp_ep(gm, node, self.rank, self.world_size, self.allreduce_strategy)
+        _insert_sharded_mxfp4_mlp_ep(
+            gm, node, self.rank, self.world_size, self.allreduce_strategy, self.dist_backend
+        )
 
 
 class FP8EPShardingInfo(EPShardingInfo, QuantizationShardingMixin):
@@ -1218,7 +1273,13 @@ def scale_names(self) -> List[str]:
 
     def apply(self, gm: GraphModule, node: Node) -> None:
         _insert_sharded_moe(
-            gm, node, self.rank, self.world_size, self.allreduce_strategy, self.scale_names()
+            gm,
+            node,
+            self.rank,
+            self.world_size,
+            self.allreduce_strategy,
+            self.dist_backend,
+            self.scale_names(),
         )
 
 
@@ -1236,7 +1297,13 @@ def scale_names(self) -> List[str]:
 
     def apply(self, gm: GraphModule, node: Node) -> None:
         _insert_sharded_moe(
-            gm, node, self.rank, self.world_size, self.allreduce_strategy, self.scale_names()
+            gm,
+            node,
+            self.rank,
+            self.world_size,
+            self.allreduce_strategy,
+            self.dist_backend,
+            self.scale_names(),
         )
 
 
@@ -1276,6 +1343,14 @@ class ShardingDim(Enum):
     BMM = "bmm"
 
 
+class DistBackend(Enum):
+    """Enum for distributed backend."""
+
+    AUTO = "auto"
+    TRTLLM = "trtllm"
+    TORCH = "torch"
+
+
 class ShardingTransformContainer(BaseModel):
     """Configuration for sharding the model."""
 
@@ -1297,6 +1372,7 @@ class ShardingTransformContainer(BaseModel):
         description="AllReduce strategy for distributed operations. "
         "Options: AUTO, NCCL, ONESHOT, TWOSHOT, MIN_LATENCY, LOWPRECISION, UB, MNNVL, NCCL_SYMMETRIC, SYMM_MEM",
     )
+    dist_backend: DistBackend = Field(default=DistBackend.AUTO)
     weight_sharding_transforms: List[WeightShardingInfo] = Field(default_factory=list)
     parameter_update_transforms: List[ParameterUpdateInfo] = Field(default_factory=list)
     bmm_transforms: List[BMMShardingInfo] = Field(default_factory=list)
@@ -1337,6 +1413,7 @@ def init_params(
         # Extract factory_source from factory_config if present
         self.factory_source = self.factory_config.get("source", ShardingConfigSource.UNKNOWN)
         self.allreduce_strategy = other.allreduce_strategy
+        self.dist_backend = other.dist_backend
         self.validate_config(ShardingSource.MANUAL)
         self.validate_config(ShardingSource.FACTORY)
 
@@ -1344,15 +1421,28 @@ def add(self, transform: ShardingTransformInfo) -> bool:
         """Append a transform only if that node was
         not sharded before. Do not overwrite existing transforms.
 
-        Automatically propagates allreduce_strategy from this config to the transform
-        if the transform doesn't already have one set.
+        Automatically propagates allreduce_strategy and dist_backend from this config
+        to the transform if the transform doesn't already have them set.
         """
-        # Inject allreduce_strategy from config into transform if it has the attribute and it's None
-        # This creates a new transform instance with the strategy set
+        # Inject allreduce_strategy and dist_backend from config into transform
+        # if they have the attributes and they're None
+        # This creates a new transform instance with the values set
+        needs_injection = False
+        transform_dict = None
+
         if hasattr(transform, "allreduce_strategy") and transform.allreduce_strategy is None:
-            # Create a new transform with the strategy injected
-            transform_dict = transform.model_dump()
+            if transform_dict is None:
+                transform_dict = transform.model_dump()
             transform_dict["allreduce_strategy"] = self.allreduce_strategy
+            needs_injection = True
+
+        if hasattr(transform, "dist_backend") and transform.dist_backend is None:
+            if transform_dict is None:
+                transform_dict = transform.model_dump()
+            transform_dict["dist_backend"] = self.dist_backend
+            needs_injection = True
+
+        if needs_injection:
             transform = type(transform)(**transform_dict)
 
         # Find the appropriate list by checking inheritance
diff --git a/tests/unittest/_torch/auto_deploy/unit/multigpu/custom_ops/test_dist.py b/tests/unittest/_torch/auto_deploy/unit/multigpu/custom_ops/test_dist.py
index 6a72868e3b5d..4bad8237f417 100644
--- a/tests/unittest/_torch/auto_deploy/unit/multigpu/custom_ops/test_dist.py
+++ b/tests/unittest/_torch/auto_deploy/unit/multigpu/custom_ops/test_dist.py
@@ -16,6 +16,7 @@ def _run_all_reduce_test(rank, world_size):
 
 def _run_all_gather_test(rank, world_size):
     x = torch.ones(10, 10).to("cuda")
+    # Test torch backend (demollm mode with Python multiprocessing)
     y = torch.ops.auto_deploy.torch_dist_all_gather(x)
 
     assert torch.sum(y) == world_size * torch.sum(x)
diff --git a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_allreduce_residual_rmsnorm_fusion.py b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_allreduce_residual_rmsnorm_fusion.py
index 83bccc33865a..408601bc6895 100644
--- a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_allreduce_residual_rmsnorm_fusion.py
+++ b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_allreduce_residual_rmsnorm_fusion.py
@@ -5,8 +5,8 @@
 from _dist_test_utils import get_device_counts
 from torch.export import export
 
+from tensorrt_llm._torch.auto_deploy.custom_ops.trtllm_dist import is_trtllm_op_available
 from tensorrt_llm._torch.auto_deploy.distributed import common as dist
-from tensorrt_llm._torch.auto_deploy.distributed.trtllm import is_trtllm_op_available
 from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm
 from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer
 from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op
@@ -32,12 +32,13 @@ def forward(self, hidden_states: torch.Tensor):
 class AllreduceResidualNorm(torch.nn.Module):
     """AllreduceResidualNorm pattern model that do residual plus x"""
 
-    def __init__(self, hidden_size, dtype):
+    def __init__(self, hidden_size, dtype, strategy):
         super().__init__()
         self.norm = RMSNorm(hidden_size, 1e-5, dtype)
+        self.strategy = strategy
 
     def forward(self, x, residual):
-        x = torch.ops.auto_deploy.torch_dist_all_reduce.default(x, "AUTO")
+        x = torch.ops.auto_deploy.trtllm_dist_all_reduce.default(x, self.strategy)
         y = residual + x
         normed = self.norm(y)
         return normed, y
@@ -46,18 +47,19 @@ def forward(self, x, residual):
 class AllreduceResidualNorm2(torch.nn.Module):
     """AllreduceResidualNorm pattern model that do x plus residual"""
 
-    def __init__(self, hidden_size, dtype):
+    def __init__(self, hidden_size, dtype, strategy):
         super().__init__()
         self.norm = RMSNorm(hidden_size, 1e-5, dtype)
+        self.strategy = strategy
 
     def forward(self, x, residual):
-        x = torch.ops.auto_deploy.torch_dist_all_reduce.default(x, "AUTO")
+        x = torch.ops.auto_deploy.trtllm_dist_all_reduce.default(x, self.strategy)
         y = x + residual
         normed = self.norm(y)
         return normed, y
 
 
-def _test_allreduce_fusion(port: int, ModuleCls):
+def _test_allreduce_fusion(port: int, ModuleCls, strategy: str):
     if not is_trtllm_op_available():
         pytest.skip("Require trtllm ops to run test_allreduce_fusion.")
 
@@ -69,7 +71,7 @@ def _test_allreduce_fusion(port: int, ModuleCls):
     residual = torch.randn(16, 16).to(dtype).cuda()
 
     # Trace the original model
-    model = ModuleCls(16, dtype)
+    model = ModuleCls(16, dtype, strategy=strategy)
     args = (
         x,
         residual,
@@ -78,10 +80,14 @@ def _test_allreduce_fusion(port: int, ModuleCls):
     # Run the original
     original_outputs, residual_original = gm(x, residual)
 
-    # Fuse ops
+    # Fuse ops with the specified strategy
     gm_transformed = InferenceOptimizer(
         None,
         {
+            "detect_sharding": {
+                "stage": "post_export",
+                "allreduce_strategy": strategy,
+            },
             "fuse_allreduce_residual_rmsnorm": {
                 "stage": "post_load_fusion",
             },
@@ -91,12 +97,21 @@ def _test_allreduce_fusion(port: int, ModuleCls):
     # Run the fused graph
     fused_outputs, residual_fused = gm_transformed(x, residual)
 
-    # Check if fused node in the graph
+    # Check if fused node in the graph and verify strategy
     has_fused_node = False
+    fused_node_strategy = None
     for node in gm_transformed.graph.nodes:
-        if is_op(node, torch.ops.dist.fused_allreduce_residual_rmsnorm):
+        if is_op(node, torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm):
             has_fused_node = True
+            # The fused node should have the strategy as the last argument
+            # args: (x, residual, weight, eps, strategy)
+            if len(node.args) >= 5:
+                fused_node_strategy = node.args[4]
+
     assert has_fused_node, "Fused node not found."
+    assert fused_node_strategy == strategy, (
+        f"Fused node strategy mismatch: expected '{strategy}', got '{fused_node_strategy}'"
+    )
 
     # Verify outputs are consistent
     assert torch.allclose(residual_original, residual_fused, atol=1e-5), (
@@ -117,11 +132,16 @@ def _test_allreduce_fusion(port: int, ModuleCls):
     [AllreduceResidualNorm, AllreduceResidualNorm2],
     ids=["residual_plus_x", "x_plus_residual"],
 )
-def test_allreduce_fusion(device_count, ModuleCls):
+@pytest.mark.parametrize(
+    "strategy",
+    ["AUTO", "NCCL", "ONESHOT"],
+    ids=["strategy_auto", "strategy_nccl", "strategy_oneshot"],
+)
+def test_allreduce_fusion(device_count, ModuleCls, strategy):
     if device_count <= 1:
         pytest.skip("Require multi GPUs to run test_allreduce_fusion.")
     port = dist.get_free_port()
 
     n_workers = device_count
     mpi_pool = MpiPoolSession(n_workers=n_workers)
-    mpi_pool.submit_sync(_test_allreduce_fusion, port=port, ModuleCls=ModuleCls)
+    mpi_pool.submit_sync(_test_allreduce_fusion, port=port, ModuleCls=ModuleCls, strategy=strategy)
diff --git a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_bmm_sharding.py b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_bmm_sharding.py
index 35a4e54e7634..ff46903aaf35 100644
--- a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_bmm_sharding.py
+++ b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_bmm_sharding.py
@@ -119,6 +119,7 @@ def _run_pattern_detection_job(
                         world_size=world_size,
                         start_idx=start_idx,
                         end_idx=end_idx,
+                        dist_backend="auto",
                     )
                 )
 
diff --git a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_dist_backend.py b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_dist_backend.py
new file mode 100644
index 000000000000..52709fe0cf61
--- /dev/null
+++ b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_dist_backend.py
@@ -0,0 +1,173 @@
+"""Test dist_backend configuration for sharding transformations."""
+
+import sys
+from pathlib import Path
+from typing import Optional
+
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Add parent directory to path for test utilities
+sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "_utils_test"))
+
+from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm
+from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer
+from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op
+
+
+class SimpleMLP(nn.Module):
+    """Simple MLP model for testing dist_backend."""
+
+    def __init__(self, in_features=32, hidden_features=64, out_features=32, bias=False):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.linear1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.linear2 = nn.Linear(hidden_features, out_features, bias=bias)
+
+    def forward(self, x):
+        y = F.relu(self.linear1(x))
+        return self.linear2(y)
+
+
+def _create_and_transform_model(
+    model: nn.Module,
+    dist_backend: Optional[str] = None,
+    world_size: int = 2,
+):
+    """Helper to create model, export to graph, and apply sharding transforms.
+
+    Args:
+        model: The model to transform
+        dist_backend: The distributed backend ('torch', 'trtllm', 'auto', or None for default)
+        world_size: World size for multi-GPU setting
+
+    Returns:
+        Transformed GraphModule
+    """
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device=device, dtype=torch.float32)
+    x = torch.randn(2, 32, device=device, dtype=torch.float32)
+
+    # Export to graph module
+    gm = torch_export_to_gm(model, args=(x,), clone=True)
+
+    # Build config
+    config = {
+        "detect_sharding": {
+            "stage": "sharding",
+        },
+        "sharding_transform_executor": {
+            "stage": "sharding",
+        },
+    }
+
+    # Add dist_backend if specified
+    if dist_backend is not None:
+        config["detect_sharding"]["dist_backend"] = dist_backend
+
+    # Run optimizer
+    optimizer = InferenceOptimizer(None, config)
+    optimizer.shared_config.local_rank = 0
+    optimizer.shared_config.world_size = world_size
+    return optimizer(None, gm)
+
+
+def _check_dist_ops(gm, expected_backend: str):
+    """Helper to check which distributed ops are present in the graph.
+
+    Args:
+        gm: GraphModule to check
+        expected_backend: Expected backend ('torch', 'trtllm', or 'any')
+
+    Returns:
+        tuple: (has_torch_ops, has_trtllm_ops)
+    """
+    has_torch_all_reduce = any(
+        is_op(n, torch.ops.auto_deploy.torch_dist_all_reduce) for n in gm.graph.nodes
+    )
+    has_trtllm_all_reduce = any(
+        is_op(n, torch.ops.auto_deploy.trtllm_dist_all_reduce) for n in gm.graph.nodes
+    )
+    has_torch_all_gather = any(
+        is_op(n, torch.ops.auto_deploy.torch_dist_all_gather) for n in gm.graph.nodes
+    )
+    has_trtllm_all_gather = any(
+        is_op(n, torch.ops.auto_deploy.trtllm_dist_all_gather) for n in gm.graph.nodes
+    )
+
+    has_torch_ops = has_torch_all_reduce or has_torch_all_gather
+    has_trtllm_ops = has_trtllm_all_reduce or has_trtllm_all_gather
+
+    if expected_backend == "torch":
+        assert has_torch_ops, (
+            f"Expected torch distributed ops when dist_backend='{expected_backend}'"
+        )
+        assert not has_trtllm_ops, (
+            f"Did not expect trtllm distributed ops when dist_backend='{expected_backend}'"
+        )
+    elif expected_backend == "trtllm":
+        assert has_trtllm_ops, (
+            f"Expected trtllm distributed ops when dist_backend='{expected_backend}'"
+        )
+        assert not has_torch_ops, (
+            f"Did not expect torch distributed ops when dist_backend='{expected_backend}'"
+        )
+    elif expected_backend == "any":
+        assert has_torch_ops or has_trtllm_ops, (
+            "Expected at least one type of distributed op in graph"
+        )
+
+    return has_torch_ops, has_trtllm_ops
+
+
+@pytest.mark.parametrize(
+    "dist_backend,expected",
+    [
+        ("torch", "torch"),
+        ("trtllm", "trtllm"),
+    ],
+)
+def test_dist_backend_explicit(dist_backend, expected):
+    """Test that explicit dist_backend forces the correct distributed ops."""
+    model = SimpleMLP()
+    gm_transformed = _create_and_transform_model(model, dist_backend=dist_backend, world_size=2)
+    _check_dist_ops(gm_transformed, expected_backend=expected)
+
+
+@pytest.mark.parametrize("dist_backend", ["auto", None])
+def test_dist_backend_auto_and_default(dist_backend):
+    """Test that dist_backend='auto' or omitting it uses auto-selection."""
+    model = SimpleMLP()
+    gm_transformed = _create_and_transform_model(model, dist_backend=dist_backend, world_size=2)
+    _check_dist_ops(gm_transformed, expected_backend="any")
+
+
+@pytest.mark.parametrize("dist_backend", ["torch", "trtllm"])
+def test_dist_backend_all_gather(dist_backend):
+    """Test dist_backend with all_gather operations (column sharding with single Linear)."""
+    model = nn.Linear(32, 64, bias=False)
+    gm_transformed = _create_and_transform_model(model, dist_backend=dist_backend, world_size=2)
+    _check_dist_ops(gm_transformed, expected_backend=dist_backend)
+
+
+if __name__ == "__main__":
+    # Run tests directly for debugging
+    print("Testing explicit backends...")
+    test_dist_backend_explicit("torch", "torch")
+    test_dist_backend_explicit("trtllm", "trtllm")
+    print("✓ Passed\n")
+
+    print("Testing auto and default backends...")
+    test_dist_backend_auto_and_default("auto")
+    test_dist_backend_auto_and_default(None)
+    print("✓ Passed\n")
+
+    print("Testing all_gather with different backends...")
+    test_dist_backend_all_gather("torch")
+    test_dist_backend_all_gather("trtllm")
+    print("✓ Passed\n")
+
+    print("\nAll tests passed!")
diff --git a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py
index 6b7c79de0518..257c2d191d42 100644
--- a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py
+++ b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_ep_sharding.py
@@ -95,6 +95,7 @@ def _run_pattern_detection_job(num_experts: int, rank: int, world_size: int) ->
                         rank=rank,
                         world_size=world_size,
                         allreduce_strategy=AllReduceStrategy.AUTO,
+                        dist_backend="auto",
                     )
                 )
             elif is_op(node, torch.ops.auto_deploy.torch_quant_fp8_moe):
@@ -104,6 +105,7 @@ def _run_pattern_detection_job(num_experts: int, rank: int, world_size: int) ->
                         rank=rank,
                         world_size=world_size,
                         allreduce_strategy=AllReduceStrategy.AUTO,
+                        dist_backend="auto",
                     )
                 )
             elif is_op(node, torch.ops.auto_deploy.torch_quant_nvfp4_moe):
@@ -113,6 +115,7 @@ def _run_pattern_detection_job(num_experts: int, rank: int, world_size: int) ->
                         rank=rank,
                         world_size=world_size,
                         allreduce_strategy=AllReduceStrategy.AUTO,
+                        dist_backend="auto",
                     )
                 )
 
diff --git a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py
index 8407580c3a3a..d34c8c0f3d20 100644
--- a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py
+++ b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py
@@ -281,6 +281,7 @@ def _run_pattern_detection_job(
                             dist_op=dist_op,
                             min_local_shape=min_local_shape,
                             allreduce_strategy=AllReduceStrategy.AUTO,
+                            dist_backend="auto",
                         )
                     )
         elif model_cls == MLP:
@@ -303,6 +304,7 @@ def _run_pattern_detection_job(
                             dist_op=dist_op,
                             min_local_shape=1,
                             allreduce_strategy=AllReduceStrategy.AUTO,
+                            dist_backend="auto",
                         )
                     )
         elif model_cls == nn.Linear:
@@ -318,6 +320,7 @@ def _run_pattern_detection_job(
                             dist_op="all_gather",
                             min_local_shape=1,
                             allreduce_strategy=AllReduceStrategy.AUTO,
+                            dist_backend="auto",
                         )
                     )
         elif model_cls == FP8MLP:
@@ -340,6 +343,7 @@ def _run_pattern_detection_job(
                             dist_op=dist_op,
                             min_local_shape=1,
                             allreduce_strategy=AllReduceStrategy.AUTO,
+                            dist_backend="auto",
                         )
                     )