diff --git a/tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py b/tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py
index 5c43d5924aa5..d603a8c27125 100644
--- a/tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py
+++ b/tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py
@@ -7,7 +7,7 @@
 import torch
 import torch.nn.functional as F
 from einops import rearrange
-from transformers import AutoConfig, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM
 
 from tensorrt_llm._torch.auto_deploy.models.patches.bamba import _bamba_mixer_torch_forward
 
@@ -190,41 +190,41 @@ def get_model_from_config_patched(config, **kwargs):
 # TODO: figure out how this can be incorporated into the export patch system
 AutoModelForCausalLM.from_config = get_model_from_config_patched
 
-_config_from_pretrained_original = AutoConfig.from_pretrained
-_nemotron_h_base_model_tp_plan = {
-    # mamba SSM layer
-    "in_proj": "mamba",
-    "out_proj": "rowwise",
-    # attention layer
-    "q_proj": "colwise",
-    "k_proj": "colwise",
-    "v_proj": "colwise",
-    "o_proj": "rowwise",
-    # NOTE: consider not sharding shared experts and/or
-    # latent projections at all, keeping them replicated.
-    # To do so, comment out the corresponding entries.
-    # moe layer: SHARED experts
-    "up_proj": "colwise",
-    "down_proj": "rowwise",
-    # MoLE: latent projections: simple shard
-    "fc1_latent_proj": "gather",
-    "fc2_latent_proj": "gather",
-}
-
-
-def get_config_from_pretrained_patched(*args, **kwargs):
-    ret = _config_from_pretrained_original(*args, **kwargs)
-    config = ret[0] if isinstance(ret, tuple) else ret
-    # heuristic to check if it's a NemotronH MoE Model
-    model_type = getattr(config, "model_type", None)
-    num_moe_layers = getattr(config, "layers_block_type", []).count("moe")
-    if model_type == "nemotron_h" and num_moe_layers > 0:
-        config.base_model_tp_plan = _nemotron_h_base_model_tp_plan
-    return (config, *ret[1:]) if isinstance(ret, tuple) else config
-
-
-# TODO: figure out how this can be incorporated into the export patch system
-AutoConfig.from_pretrained = get_config_from_pretrained_patched
+# _config_from_pretrained_original = AutoConfig.from_pretrained
+# _nemotron_h_base_model_tp_plan = {
+#     # mamba SSM layer
+#     "in_proj": "mamba",
+#     "out_proj": "rowwise",
+#     # attention layer
+#     "q_proj": "colwise",
+#     "k_proj": "colwise",
+#     "v_proj": "colwise",
+#     "o_proj": "rowwise",
+#     # NOTE: consider not sharding shared experts and/or
+#     # latent projections at all, keeping them replicated.
+#     # To do so, comment out the corresponding entries.
+#     # moe layer: SHARED experts
+#     "up_proj": "colwise",
+#     "down_proj": "rowwise",
+#     # MoLE: latent projections: simple shard
+#     "fc1_latent_proj": "gather",
+#     "fc2_latent_proj": "gather",
+# }
+
+
+# def get_config_from_pretrained_patched(*args, **kwargs):
+#     ret = _config_from_pretrained_original(*args, **kwargs)
+#     config = ret[0] if isinstance(ret, tuple) else ret
+#     # heuristic to check if it's a NemotronH MoE Model
+#     model_type = getattr(config, "model_type", None)
+#     num_moe_layers = getattr(config, "layers_block_type", []).count("moe")
+#     if model_type == "nemotron_h" and num_moe_layers > 0:
+#         config.base_model_tp_plan = _nemotron_h_base_model_tp_plan
+#     return (config, *ret[1:]) if isinstance(ret, tuple) else config
+
+
+# # TODO: figure out how this can be incorporated into the export patch system
+# AutoConfig.from_pretrained = get_config_from_pretrained_patched
 
 # TODO: figure out how this can be incorporated into the export patch system
 # Only patch if the module isn't available
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
index 6230cd3031d9..02f7e226c173 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
@@ -16,10 +16,8 @@
        happens automatically via the checkpoint loading hook added in step 2c.
 """
 
-import operator
 import re
-from collections import defaultdict
-from typing import Any, DefaultDict, Dict, List, Set, Tuple, Type
+from typing import Any, Dict, List, Tuple, Type, Union
 
 import torch
 from pydantic import Field, field_validator
@@ -33,8 +31,11 @@
     bfs,
     extract_weight_node,
     filtered_nodes,
-    identify_regions_between_residuals,
+    get_all_layer_subgraphs,
+    is_any_attention_op,
     is_any_lin_op,
+    is_any_moe_op,
+    is_any_ssm_op,
     is_op,
     subgraph,
 )
@@ -78,7 +79,7 @@ class ShardingTransformConfig(TransformConfig):
         ]
     )
     sharding_dims: List[ShardingDim] = Field(
-        default_factory=lambda: [ShardingDim.SSM, ShardingDim.TP, ShardingDim.EP, ShardingDim.BMM]
+        default_factory=lambda: [ShardingDim.TP, ShardingDim.EP, ShardingDim.BMM]
     )
     allreduce_strategy: AllReduceStrategy = Field(
         default=AllReduceStrategy.AUTO,
@@ -152,33 +153,38 @@ def check_and_apply(transform: ShardingTransformInfo) -> bool:
             is_clean=num_matches == 0,
             has_valid_shapes=num_matches == 0,
         )
-        # exit()
         return gm, info
 
 
 def _process_simple_shard(
-    nodes_linear: Dict[Node, List[Node]],
+    nodes_linear: Union[Dict[Node, List[Node]], List[Node]],
     rank: int,
     world_size: int,
-    sharding_config: ShardingTransformContainer,
+    transform_container: ShardingTransformContainer,
+    layer_type: LayerType = LayerType.MLP,
 ) -> int:
     # for every linear node:
     # --> row_split (dim 0 of weight) + all_gather (dim -1 of output)
+    # if nodes_linear is a dict, flatten it to a 1D list of nodes
+
+    if isinstance(nodes_linear, dict):
+        nodes_linear = [n for group in nodes_linear.values() for n in group]
+
     num_simple_shards = 0
-    for node_group in nodes_linear.values():
-        for n in node_group:
-            num_simple_shards += int(
-                sharding_config.add(
-                    WeightShardingInfo.from_node(
-                        n,
-                        split_dim=SplitDimension.COLUMN,
-                        rank=rank,
-                        world_size=world_size,
-                        dist_op="all_gather",
-                        min_local_shape=1,
-                    )
+    for n in nodes_linear:
+        num_simple_shards += int(
+            transform_container.add(
+                WeightShardingInfo.from_node(
+                    n,
+                    split_dim=SplitDimension.COLUMN,
+                    rank=rank,
+                    world_size=world_size,
+                    dist_op="all_gather",
+                    min_local_shape=1,
+                    layer_type=layer_type,
                 )
             )
+        )
     return num_simple_shards
 
 
@@ -246,9 +252,6 @@ def _apply(
                 ad_logger.info(
                     f"Running autodeploy sharding heuristics: {transform_container.sharding_dims}"
                 )
-                if ShardingDim.SSM in transform_container.sharding_dims:
-                    info += detect_ssm_shard(gm, transform_container)
-
                 # run TP sharding across ranks
                 if ShardingDim.TP in transform_container.sharding_dims:
                     info += detect_column_row_shard(gm, transform_container)
@@ -267,7 +270,7 @@ def _apply(
 def _process_ssm_sharding(
     gm: GraphModule,
     entry_node: Node,
-    sharding_config: ShardingTransformContainer,
+    transform_container: ShardingTransformContainer,
     rank: int,
     world_size: int,
     min_local_shape: int = 1,
@@ -313,6 +316,24 @@ def _process_ssm_sharding(
         "conv1d": split_sizes_1,
     }
 
+    # ##############################################################
+    # ####### shard the entry_node (the first linear layer) ########
+    # ##############################################################
+    if not transform_container.add(
+        WeightShardingInfo.from_node(
+            entry_node,
+            split_dim=SplitDimension.COLUMN,
+            rank=rank,
+            world_size=world_size,
+            dist_op=None,
+            min_local_shape=min_local_shape,
+            fused_weight_dims=fused_weight_dims["in_proj"],
+            layer_type=LayerType.MAMBA,
+        )
+    ):
+        # the layer was already sharded. Skipping.
+        return 0
+
     # # ##############################################################
     # # ############## update split nodes ############################
     # # ##############################################################
@@ -320,7 +341,7 @@ def _process_ssm_sharding(
     split_args_0[1] = [s // world_size for s in split_args_0[1]]
     split_args_1 = list(split_nodes[1].args)
     split_args_1[1] = [s // world_size for s in split_args_1[1]]
-    sharding_config.add(
+    transform_container.add(
         ParameterUpdateInfo(
             rank=rank,
             world_size=world_size,
@@ -328,7 +349,7 @@ def _process_ssm_sharding(
             args=tuple(split_args_0),
         )
     )
-    sharding_config.add(
+    transform_container.add(
         ParameterUpdateInfo(
             rank=rank,
             world_size=world_size,
@@ -349,27 +370,12 @@ def _process_ssm_sharding(
     # This one is also sharded, so we need to update this parameter
     conv_args = list(conv1d_node.args)
     conv_args[-1] = conv1d_node.args[-1] // world_size
-    sharding_config.add(
+    transform_container.add(
         ParameterUpdateInfo(
             rank=rank, world_size=world_size, target_node=conv1d_node.name, args=tuple(conv_args)
         )
     )
 
-    # ##############################################################
-    # ####### shard the entry_node (the first linear layer) ########
-    # ##############################################################
-    sharding_config.add(
-        WeightShardingInfo.from_node(
-            entry_node,
-            split_dim=SplitDimension.COLUMN,
-            rank=rank,
-            world_size=world_size,
-            dist_op=None,
-            min_local_shape=min_local_shape,
-            fused_weight_dims=fused_weight_dims["in_proj"],
-        )
-    )
-
     # ##############################################################
     # ############## shard the remaining weights ###################
     # ##############################################################
@@ -396,7 +402,7 @@ def _process_ssm_sharding(
                 break
 
         # Shard the weight tensor (also updates the parameter in the module)
-        sharding_config.add(
+        transform_container.add(
             WeightShardingInfo.from_node(
                 list(weight_node.users)[0],
                 split_dim=SplitDimension.COLUMN,
@@ -405,6 +411,7 @@ def _process_ssm_sharding(
                 dist_op=None,
                 min_local_shape=min_local_shape,
                 fused_weight_dims=fused_dims,
+                layer_type=LayerType.MAMBA,
             )
         )
 
@@ -424,7 +431,7 @@ def _process_ssm_sharding(
             args = list(view_node.args)
             view_shape[2] = view_shape[2] // world_size
             args[1] = tuple(view_shape)
-            sharding_config.add(
+            transform_container.add(
                 ParameterUpdateInfo(
                     rank=rank, world_size=world_size, target_node=view_node.name, args=tuple(args)
                 )
@@ -434,32 +441,73 @@ def _process_ssm_sharding(
     ##############################################################
     ############## shard the out_proj node #######################
     ##############################################################
-    sharding_config.add(
+    transform_container.add(
         WeightShardingInfo.from_node(
             out_proj_node,
             split_dim=SplitDimension.ROW,
             rank=rank,
             world_size=world_size,
             dist_op="all_reduce",
+            layer_type=LayerType.MAMBA,
         )
     )
     return 1
 
 
 def _process_column_sharding(
-    gm: GraphModule,
     linear_nodes: List[Node],
-    sharding_config: ShardingTransformContainer,
+    subgraph_nodes: Union[List[Node], None],
+    transform_container: ShardingTransformContainer,
     rank: int,
     world_size: int,
     min_local_shape: int = 1,
-    fused_weight: bool = False,
 ) -> None:
     """
     Parse the column sharding from the candidate nodes and update the view and split nodes accordingly.
     """
+    if subgraph_nodes is None:
+        subgraph_nodes = subgraph(linear_nodes, boundary_condition=is_any_lin_op)
+    fused_weight_dims = None
+    # check if there are split nodes in the subgraph. They may indicate fused weights (e.g., QKV)
+    split_nodes = list(filtered_nodes(subgraph_nodes, ops=[torch.ops.aten.split_with_sizes]))
+    if len(split_nodes) > 0:
+        assert len(linear_nodes) == 1
+        linear_node = linear_nodes[0]
+        assert len(split_nodes) == 1, "Expecting exactly one split node for fused weights"
+        fused_weight_dims = split_nodes[0].args[1]
+    slice_nodes = list(filtered_nodes(subgraph_nodes, ops=[torch.ops.aten.slice]))
+    if len(slice_nodes) > 0:
+        # we are probably in fused QKV case with single linear node and 3 slice nodes
+        assert len(linear_nodes) == 1
+        linear_node = linear_nodes[0]
+        assert all(
+            s.args[1] == 2 for s in filtered_nodes(linear_node.users, ops=torch.ops.aten.slice)
+        ), "Expecting slice nodes to slice tensor over dim=2"
+        fused_weight_dims = [s.args[3] - s.args[2] for s in linear_node.users]
+        weight_dim = linear_node.meta["val"].shape[2]
+        if sum(fused_weight_dims) != weight_dim:
+            if fused_weight_dims[-1] > weight_dim:
+                fused_weight_dims[-1] = weight_dim - sum(fused_weight_dims[:-1])
+            else:
+                ad_logger.warning(
+                    f"Fused weight dims {fused_weight_dims} do not sum to weight dim {weight_dim}. Skipping."
+                )
+                return
+    chunk_nodes = list(filtered_nodes(subgraph_nodes, ops=torch.ops.aten.chunk))
+    if len(chunk_nodes) > 0:
+        assert len(linear_nodes) == 1
+        linear_node = linear_nodes[0]
+        assert len(chunk_nodes) == 1, "Expecting exactly one chunk node for fused weights"
+        num_chunks = chunk_nodes[0].args[1]
+        weight_dim = linear_node.meta["val"].shape[2]
+        fused_weight_dims = [weight_dim // num_chunks] * num_chunks
+
+    # check if there are any attention nodes in the subgraph
+    attention_nodes = list(filtered_nodes(subgraph_nodes, is_any_attention_op))
+
+    added_nodes: int = 0
     for linear_node in linear_nodes:
-        sharding_config.add(
+        added_nodes += transform_container.add(
             WeightShardingInfo.from_node(
                 linear_node,
                 split_dim=SplitDimension.COLUMN,
@@ -467,15 +515,13 @@ def _process_column_sharding(
                 world_size=world_size,
                 dist_op=None,  # for column sharding, no dist op is performed
                 min_local_shape=min_local_shape,
+                fused_weight_dims=fused_weight_dims,
+                layer_type=LayerType.ATTENTION if len(attention_nodes) > 0 else LayerType.MLP,
             )
         )
-
-    # get the subgraph of this module. Subgraph boundary is the next linear node.
-    next_lin_node, _ = bfs(linear_nodes[0], is_any_lin_op, include_root=False)
-    subgraph_nodes = subgraph(
-        linear_nodes,
-        [next_lin_node],
-    )
+    if added_nodes == 0:
+        ad_logger.debug("No nodes were added for column sharding. Skipping.")
+        return
 
     nodes_to_validate = [
         n for n in subgraph_nodes if is_op(n, [torch.ops.aten.view, torch.ops.aten.reshape])
@@ -488,9 +534,9 @@ def _process_column_sharding(
             continue
         if len(view_shape) >= 3 and isinstance(view_shape[2], int) and view_shape[2] != -1:
             args = list(view_node.args)
-            view_shape[2] = view_shape[2] // world_size
+            view_shape[2] = -1
             args[1] = tuple(view_shape)
-            sharding_config.add(
+            transform_container.add(
                 ParameterUpdateInfo(
                     rank=rank, world_size=world_size, target_node=view_node.name, args=tuple(args)
                 )
@@ -498,25 +544,41 @@ def _process_column_sharding(
             ad_logger.debug(f"\nUpdated view node {view_node} arguments to {view_node.args}")
 
     # if fused_weight_dims is provided, we need to update all split sizes
-    if fused_weight:
-        assert len(linear_nodes) == 1, "Fused weight should be only one linear node"
-        node = linear_nodes[0]
+    if fused_weight_dims is not None:
         assert world_size is not None, "World size is required to update the split node params"
-        assert len(node.users) == 1, "Fused linear node should have only one user: a split node"
-        user = list(node.users)[0]
-        if is_op(user, [torch.ops.aten.split_with_sizes]):
+
+        # fused weight may either be processed by several slice nodes or a single split node
+        assert len(split_nodes) > 0 or len(slice_nodes) > 0 or len(chunk_nodes) > 0, (
+            "Expecting at least one split or slice or chunk node for fused weights"
+        )
+
+        assert len(linear_nodes) == 1, "Expecting exactly one linear node for fused weights"
+        linear_node = linear_nodes[0]
+        if len(split_nodes) > 0:
+            user = split_nodes[0]
             orig_sizes = user.args[1]
             new_sizes = [orig_sizes[i] // world_size for i in range(len(orig_sizes))]
             args = list(user.args)
             args[1] = new_sizes
-            sharding_config.add(
+            transform_container.add(
                 ParameterUpdateInfo(
                     rank=rank, world_size=world_size, target_node=user.name, args=tuple(args)
                 )
             )
-            ad_logger.debug(
-                f"\nInserted parameter update transformation for split node {user} arguments to {user.args}"
-            )
+        elif len(slice_nodes) > 0:
+            for slice_node in filtered_nodes(linear_node.users, ops=torch.ops.aten.slice):
+                args = list(slice_node.args)
+                args[2] = args[2] // world_size
+                args[3] = args[3] // world_size
+                transform_container.add(
+                    ParameterUpdateInfo(
+                        rank=rank,
+                        world_size=world_size,
+                        target_node=slice_node.name,
+                        args=tuple(args),
+                    )
+                )
+        # chunk nodes do not need to be updated
 
 
 def detect_sharding_from_config(
@@ -529,7 +591,7 @@ def detect_sharding_from_config(
     TODO: currently, it applies only to TP sharding.
     Args:
         gm: Graph module to apply transformations to
-        sharding_config: Predefined sharding configuration
+        transform_container: containing predefined sharding configuration
     """
     # check if config is valid.
     # 1. it is a Dict[str, str]
@@ -575,6 +637,8 @@ def detect_sharding_from_config(
     num_shards = 0
     num_simple_shards = 0
     num_row_col_shards = 0
+    num_attention_shards = 0
+    num_ssm_shards = 0
 
     for lin_node in filtered_nodes(gm.graph.nodes, is_any_lin_op):
         # use node's weight name to get the module name
@@ -582,8 +646,10 @@ def detect_sharding_from_config(
 
         if any(attn_name in module_name for attn_name in attn_names):
             min_local_shape = head_dim
+            layer_type = LayerType.ATTENTION
         else:
             min_local_shape = 1
+            layer_type = LayerType.MLP
 
         # use regex to find if module_name matches any of the keys in sharding_config
         for key in tp_plan.keys():
@@ -594,21 +660,17 @@ def detect_sharding_from_config(
             pattern_string = pattern_string.replace("*", "@")
             pattern_regex = re.escape(pattern_string).replace("@", ".*")
             if re.match(pattern_regex, module_name):
-                num_shards += 1
                 # we have a match. Get the config for this layer
                 config = tp_plan[key]
                 if config == "colwise":
-                    if transform_container.add(
-                        WeightShardingInfo.from_node(
-                            lin_node,
-                            split_dim=SplitDimension.COLUMN,
-                            rank=rank,
-                            world_size=world_size,
-                            dist_op=None,
-                            min_local_shape=min_local_shape,
-                        )
-                    ):
-                        num_row_col_shards += 1
+                    _process_column_sharding(
+                        linear_nodes=[lin_node],
+                        subgraph_nodes=None,
+                        transform_container=transform_container,
+                        rank=rank,
+                        world_size=world_size,
+                        min_local_shape=min_local_shape,
+                    )
                 elif config == "rowwise":
                     if transform_container.add(
                         WeightShardingInfo.from_node(
@@ -618,22 +680,20 @@ def detect_sharding_from_config(
                             world_size=world_size,
                             dist_op="all_reduce",
                             min_local_shape=min_local_shape,
+                            layer_type=layer_type,
                         )
                     ):
+                        if layer_type == LayerType.ATTENTION:
+                            num_attention_shards += 1
                         num_row_col_shards += 1
                 elif config == "mamba":
-                    transform_container.add(
-                        WeightShardingInfo.from_node(
-                            lin_node,
-                            split_dim=SplitDimension.COLUMN,
-                            rank=rank,
-                            world_size=world_size,
-                            dist_op=None,
-                            min_local_shape=min_local_shape,
-                            layer_type=LayerType.MAMBA,
-                        )
-                    )
-                    num_row_col_shards += 1
+                    if (
+                        _process_ssm_sharding(gm, lin_node, transform_container, rank, world_size)
+                        > 0
+                    ):
+                        num_ssm_shards += 1
+                        num_row_col_shards += 1
+
                 elif "sequence" in config:
                     # TODO: Sequence parallelism is not supported yet.
                     ad_logger.warning("Sequence parallelism is not supported yet. Skipping.")
@@ -651,6 +711,7 @@ def detect_sharding_from_config(
                                     world_size=world_size,
                                     dist_op=None,
                                     min_local_shape=min_local_shape,
+                                    layer_type=layer_type,
                                 )
                             )
                         elif col_row_action == "rowwise":
@@ -662,6 +723,7 @@ def detect_sharding_from_config(
                                     world_size=world_size,
                                     dist_op="all_reduce",
                                     min_local_shape=min_local_shape,
+                                    layer_type=layer_type,
                                 )
                             ):
                                 num_row_col_shards += 1
@@ -681,29 +743,22 @@ def detect_sharding_from_config(
                             world_size=world_size,
                             dist_op="all_gather",
                             min_local_shape=1,
+                            layer_type=layer_type,
                         )
                     ):
                         num_simple_shards += 1
                 else:
-                    ad_logger.warning(
-                        f"Unsupported sharding action {config}. Fallback to simple shard"
-                    )
-                    transform_container.add(
-                        WeightShardingInfo.from_node(
-                            lin_node,
-                            split_dim=SplitDimension.COLUMN,
-                            rank=rank,
-                            world_size=world_size,
-                            dist_op="all_gather",
-                            min_local_shape=1,
-                        )
+                    ad_logger.debug(
+                        f"Unsupported sharding action {config}. "
+                        f"Linear node {lin_node} will not be sharded."
                     )
                 # after successful match, break the loop
                 break
 
+    num_shards = num_simple_shards + num_row_col_shards
     ad_logger.info(
-        f"Applied {num_shards} TP shards (simple: {num_simple_shards}, "
-        f"row-col pattern: {num_row_col_shards})"
+        f"Applied {num_shards} TP shards from config. Simple: {num_simple_shards}, "
+        f"row-col: {num_row_col_shards} (including: ssm: {num_ssm_shards}, attention: {num_attention_shards})"
     )
 
     num_matches = len(transform_container.weight_sharding_transforms)
@@ -718,7 +773,7 @@ def detect_sharding_from_config(
 
 def detect_ssm_shard(
     gm: GraphModule,
-    sharding_config: ShardingTransformContainer,
+    transform_container: ShardingTransformContainer,
 ) -> TransformInfo:
     """A transformation to apply sharding to the model following SSM parallelism.
     TODO: This is a TEMPORARY place for this logic due to the incompatibility between the
@@ -726,14 +781,14 @@ def detect_ssm_shard(
     The goal is to have a unified single pass over the graph to detect layers and apply
     appropriate sharding transformations.
     """
-    rank, world_size = sharding_config.rank, sharding_config.world_size
+    rank, world_size = transform_container.rank, transform_container.world_size
     if world_size < 2:
         ad_logger.info("Skipping TP sharding for single device")
         return TransformInfo(skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True)
     ad_logger.info("Running SSM sharding detection")
 
     # find all ssm nodes in the graph
-    ssm_nodes = filtered_nodes(gm.graph.nodes, ops=torch.ops.auto_deploy.torch_ssm)
+    ssm_nodes = filtered_nodes(gm.graph.nodes, is_any_ssm_op)
     num_ssm_shards = 0
     for ssm_node in ssm_nodes:
         # We assume that one ssm node defines a subgraph corresponding
@@ -742,7 +797,7 @@ def detect_ssm_shard(
         in_proj_node, _ = bfs(ssm_node, is_any_lin_op, attr_next="args", include_root=False)
 
         num_ssm_shards += int(
-            _process_ssm_sharding(gm, in_proj_node, sharding_config, rank, world_size)
+            _process_ssm_sharding(gm, in_proj_node, transform_container, rank, world_size)
         )
 
     ad_logger.info(f"Found {num_ssm_shards} SSM shards")
@@ -753,7 +808,7 @@ def detect_ssm_shard(
 
 def detect_column_row_shard(
     gm: GraphModule,
-    sharding_config: ShardingTransformContainer,
+    transfrom_container: ShardingTransformContainer,
 ) -> TransformInfo:
     """A transformation to apply sharding to the model following tensor parallelism.
 
@@ -772,193 +827,128 @@ def detect_column_row_shard(
     splitting, e.g., the individual heads into smaller shards.
     """
     ad_logger.debug("Before sharding graph: " + str(gm))
-
-    rank, world_size = sharding_config.rank, sharding_config.world_size
-    if world_size < 2:
-        ad_logger.info("Skipping TP sharding for single device")
-        return TransformInfo(skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True)
+    rank, world_size = transfrom_container.rank, transfrom_container.world_size
 
     assert isinstance(gm, GraphModule), "Expecting GraphModule"
-
     ad_logger.info("Running TP sharding detection")
+    linear_nodes = list(filtered_nodes(gm.graph.nodes, is_any_lin_op))
+    if len(linear_nodes) == 0:
+        ad_logger.warning("Could not find any linear nodes in the graph. Skipping TP sharding.")
+        return TransformInfo(skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True)
 
-    # find boundary nodes of regions we want to shard
-    boundary_nodes = identify_regions_between_residuals(gm)
-
-    # TODO: continue updating these lists
-    # pointwise ops that don't affect the sharder
-    pointwise_ops = {
-        torch.ops.aten.gelu,
-        torch.ops.aten.leaky_relu,
-        torch.ops.aten.mul,
-        torch.ops.aten.relu,
-        torch.ops.aten.sigmoid,
-        torch.ops.aten.silu,
-        torch.ops.aten.tanh,
-        torch.ops.aten.contiguous,
-    }
-
-    # acceptable attention nodes between sharded GEMMs
-    shardable_attention_nodes = {
-        torch.ops.auto_deploy.torch_attention_sdpa,
-        torch.ops.auto_deploy.torch_attention,
-    }
-
-    # This is a heuristic. Basically, we assume those are okay to shard if we also encounter an
-    # attention node because we know that those ops must be compatible with the attention op. Now
-    # since the attention op is shardable, we will assume those are as well if used in conjunction
-    # with the attention op.
-    shardable_nodes_with_attention = {
-        torch.ops.aten.view,
-        torch.ops.aten.reshape,
-        torch.ops.auto_deploy.flashinfer_rope,
-        operator.getitem,
-    }
+    layer_subgraphs, unprocessed_linear_nodes = get_all_layer_subgraphs(gm)
 
-    # let's look at linear nodes we can identify between pairs of boundary nodes
-    # There is three potential cases we can handle:
-    # 1. No linear nodes:
-    #       --> just continue
-    # 2. Two groups of linear nodes and we can account for all to the view nodes:
-    #       --> row_split (dim 0) 1st group + check for supported nodes +
-    #           col_split (dim 1) 2nd group + all_reduce output of 2nd group
-    # 3. Linear nodes that are not in two groups or we cannot account for all nodes:
-    #       --> row_split (dim 0 of weight) + all_gather (dim -1 of output) output
     num_shards = 0
     num_simple_shards = 0
-    num_row_col_shards = 0
-    for n_start, n_end in zip(boundary_nodes[:-1], boundary_nodes[1:]):
-        # we iterate through all nodes between the two boundary nodes and store linear nodes
-        # sorted by their input activation node. We also store remaining nodes.
-        nodes_linear: DefaultDict[Node, List[Node]] = defaultdict(list)
-        attention_nodes: Set[Node] = set()
-        attention_related_nodes: Set[Node] = set()
-        unaccounted_nodes: Set[Node] = set()
-        current_node = n_start
-        while current_node != n_end:
-            if is_any_lin_op(current_node):
-                nodes_linear[current_node.args[0]].append(current_node)
-            elif is_op(current_node, shardable_attention_nodes):
-                attention_nodes.add(current_node)
-            elif is_op(current_node, shardable_nodes_with_attention):
-                attention_related_nodes.add(current_node)
-            elif not is_op(current_node, pointwise_ops):
-                unaccounted_nodes.add(current_node)
-            current_node = current_node.next
-            assert current_node, "Could not identify next node"
-
-        # nothing to shard
-        if len(nodes_linear) == 0:
-            continue
-
+    num_ssm_shards = 0
+    num_attention_shards = 0
+    num_column_row_shards = 0
+    for opening, layer_subgraph, closing in layer_subgraphs:
+        nodes_linear = opening + [closing]
         num_shards += 1
 
-        if sharding_config.simple_shard_only:
-            ad_logger.debug(f"Forcing Simple Shard: Linear groups: {nodes_linear}")
-            num_simple_shards += _process_simple_shard(
-                nodes_linear, rank, world_size, sharding_config
-            )
-            continue
+        ssm_nodes = list(filtered_nodes(layer_subgraph, is_any_ssm_op))
+        attention_nodes = list(filtered_nodes(layer_subgraph, is_any_attention_op))
+        min_local_shape = 1
+        layer_type = (
+            LayerType.MAMBA
+            if len(ssm_nodes) > 0
+            else LayerType.ATTENTION
+            if len(attention_nodes) > 0
+            else LayerType.MLP
+        )
 
-        # simple shard when we have != 2 groups of linear nodes
-        if len(nodes_linear) != 2:
-            ad_logger.debug(f"Linear groups: {nodes_linear}")
+        if transfrom_container.simple_shard_only:
+            ad_logger.debug(
+                f"Forcing Simple Shard on nodes: {nodes_linear} with layer type: {layer_type}"
+            )
             num_simple_shards += _process_simple_shard(
-                nodes_linear, rank, world_size, sharding_config
+                nodes_linear, rank, world_size, transfrom_container, layer_type=layer_type
             )
             continue
 
-        # let's look at the unnacounted nodes. They are okay as long as they fall before the
-        # first linear node or after the last linear node, i.e., outside the sharded region
-        lin_nodes_flat: Set[Node] = {n for group in nodes_linear.values() for n in group}
-        lin_nodes_passed: Set[Node] = set()
-        current_node = n_start
-        while current_node != n_end:
-            # check if this is another linear node
-            if current_node in lin_nodes_flat:
-                lin_nodes_passed.add(current_node)
-
-            # check if we are OUTSIDE sharded region
-            if len(lin_nodes_passed) == 0 or lin_nodes_passed == lin_nodes_flat:
-                # remove node from unaccounted nodes since we are outside and it doesn't matter
-                unaccounted_nodes.discard(current_node)
-                attention_related_nodes.discard(current_node)
-                attention_nodes.discard(current_node)
-
-            current_node = current_node.next
-
-        # let's post-process the attention-related nodes
-        # we can disregard them if we also see attention nodes and we assume they are compatible
-        if len(attention_nodes) > 0:
-            attention_related_nodes.clear()
-
-        # check if any unaccounted nodes are left. If so, do a simply shard
-        if unaccounted_nodes or attention_related_nodes:
-            ad_logger.debug(f"Unaccounted nodes: {unaccounted_nodes}")
-            num_simple_shards += _process_simple_shard(
-                nodes_linear, rank, world_size, sharding_config
+        if len(ssm_nodes) > 0:
+            # Mamba layers need special handling due to the fused weights for in_proj and conv1d
+            assert len(ssm_nodes) == 1, "Expected exactly one SSM node in layer subgraph"
+            assert len(opening) == 1, "Expected exactly one opening node in Mamba layer"
+            ad_logger.debug(f"Found SSM nodes in layer subgraph: {ssm_nodes}")
+            num_ssm_shards += _process_ssm_sharding(
+                gm, opening[0], transfrom_container, rank, world_size
             )
             continue
 
-        # If we can account for all sharded nodes, we can do a two-way shard
-        # --> row_split (dim 0) + col_split (dim 1) + all_reduce
-
-        # check if we are sharding the attention block
-        if attention_nodes:
+        if len(attention_nodes) > 0:
+            ad_logger.debug(f"Found attention nodes in layer subgraph: {attention_nodes}")
             if len(attention_nodes) > 1:
                 # Column-row shard boundary region detection is probably wrong - there should be
                 # only one attention operation. Fall back to simple shard.
-                ad_logger.debug(f"More than one attention node: {unaccounted_nodes}")
+                ad_logger.debug(f"More than one attention node: {attention_nodes}")
                 num_simple_shards += _process_simple_shard(
-                    nodes_linear, rank, world_size, sharding_config
+                    nodes_linear, rank, world_size, transfrom_container, layer_type=layer_type
                 )
                 continue
             # Extract head dimension. We cannot shard below the head_dim size.
             # Assume that head_dim is the last (innermost) dimension of the tensor
             min_local_shape = attention_nodes.pop().meta["val"].shape[-1]
-        else:
-            min_local_shape = 1
-
-        # We are inserting column-row shard for each group of linear enodes
-        # This may require parameter update of nodes whose args depend on (sharded) dimensions,
-        # such as view or split nodes.
-        nodes_to_column_shard = list(nodes_linear.values())[0]
-        nodes_to_row_shard = list(nodes_linear.values())[1]
-        if len(nodes_to_row_shard) != 1:
-            ad_logger.warning(
-                "Expecting only one linear node for row sharding, but got %s",
-                len(nodes_to_row_shard),
-            )
-            num_simple_shards += _process_simple_shard(
-                nodes_linear, rank, world_size, sharding_config
-            )
-            continue
+            # if the QKV projection is fused, check if num_kv_heads is divisible by world_size
+            if len(opening) == 1:
+                qkv_proj_node = opening[0]
+                slice_nodes = list(filtered_nodes(qkv_proj_node.users, ops=torch.ops.aten.slice))
+                if len(slice_nodes) > 0:
+                    # extract num_kv_heads * head_dim from the second slice node
+                    assert len(slice_nodes) == 3, "Expecting exactly 3 slice nodes for fused QKV"
+                    num_kv_heads = (
+                        slice_nodes[1].args[3] - slice_nodes[1].args[2]
+                    ) // min_local_shape
+                    if num_kv_heads % world_size != 0:
+                        ad_logger.debug(
+                            f"num_kv_heads {num_kv_heads} is not divisible by world_size {world_size}. "
+                            f"Falling back to simple shard."
+                        )
+                        num_simple_shards += _process_simple_shard(
+                            nodes_linear,
+                            rank,
+                            world_size,
+                            transfrom_container,
+                            layer_type=layer_type,
+                        )
+                        # TODO: handle the case where num_kv_heads is not divisible by world_size
+                        continue
 
         # column-row sharding
         _process_column_sharding(
-            gm,
-            linear_nodes=nodes_to_column_shard,
-            sharding_config=sharding_config,
+            linear_nodes=opening,
+            subgraph_nodes=layer_subgraph,
+            transform_container=transfrom_container,
             rank=rank,
             world_size=world_size,
             min_local_shape=min_local_shape,
         )
 
         # shard single row node
-        if sharding_config.add(
+        if transfrom_container.add(
             WeightShardingInfo.from_node(
-                nodes_to_row_shard[0],
+                closing,
                 split_dim=SplitDimension.ROW,
                 rank=rank,
                 world_size=world_size,
                 dist_op="all_reduce",
                 min_local_shape=min_local_shape,
+                layer_type=layer_type,
             )
         ):
-            num_row_col_shards += 1
+            num_column_row_shards += 1
+            if layer_type == LayerType.ATTENTION:
+                num_attention_shards += 1
 
+    # simple shard remaining linear nodes
+    num_simple_shards += _process_simple_shard(
+        unprocessed_linear_nodes, rank, world_size, transfrom_container
+    )
+    num_column_row_shards += num_ssm_shards
     ad_logger.info(
-        f"Found {num_shards} TP shards (simple: {num_simple_shards}, row-col: {num_row_col_shards})"
+        f"Heuristics found {num_shards} TP shards. Simple: {num_simple_shards}, "
+        f"row-col: {num_column_row_shards} (including: ssm: {num_ssm_shards}, attention: {num_attention_shards})"
     )
     return TransformInfo(
         skipped=False, num_matches=num_shards, is_clean=False, has_valid_shapes=False
@@ -966,7 +956,7 @@ def detect_column_row_shard(
 
 
 def detect_dp_bmm_shard(
-    gm: GraphModule, sharding_config: ShardingTransformContainer
+    gm: GraphModule, transform_container: ShardingTransformContainer
 ) -> TransformInfo:
     """A transformation to apply sharding to batched matrix multiplications in the graph.
 
@@ -977,7 +967,7 @@ def detect_dp_bmm_shard(
     We'll also assume that the inputs to BMM are broadcasted across the devices already.
     """
     ad_logger.debug("Before sharding graph: " + str(gm))
-    rank, world_size = sharding_config.rank, sharding_config.world_size
+    rank, world_size = transform_container.rank, transform_container.world_size
     if world_size < 2:
         ad_logger.info("Skipping DP BMM sharding for single device")
         return TransformInfo(skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True)
@@ -1024,7 +1014,7 @@ def detect_dp_bmm_shard(
             start_idx = remainder + rank * base_size
             end_idx = start_idx + base_size
 
-        sharding_config.add(
+        transform_container.add(
             BMMShardingInfo(
                 target_node=node.name,
                 rank=rank,
@@ -1047,10 +1037,12 @@ def detect_dp_bmm_shard(
     )
 
 
-def detect_ep_shard(gm: GraphModule, sharding_config: ShardingTransformContainer) -> TransformInfo:
+def detect_ep_shard(
+    gm: GraphModule, transform_container: ShardingTransformContainer
+) -> TransformInfo:
     ad_logger.debug("Before sharding graph: " + str(gm))
 
-    rank, world_size = sharding_config.rank, sharding_config.world_size
+    rank, world_size = transform_container.rank, transform_container.world_size
     if world_size < 2:
         ad_logger.info("Skipping EP sharding for single device")
         return TransformInfo(skipped=True, num_matches=0, is_clean=True, has_valid_shapes=True)
@@ -1058,17 +1050,9 @@ def detect_ep_shard(gm: GraphModule, sharding_config: ShardingTransformContainer
     assert isinstance(gm, GraphModule), "Expecting GraphModule"
     num_moe_patterns = 0
     for node in list(gm.graph.nodes):
-        if not is_op(
-            node,
-            (
-                torch.ops.auto_deploy.torch_moe,
-                torch.ops.auto_deploy.torch_quant_fp8_moe,
-                torch.ops.auto_deploy.torch_quant_nvfp4_moe,
-                torch.ops.auto_deploy.triton_mxfp4_moe,
-            ),
-        ):
+        if not is_any_moe_op(node):
             continue
-        if sharding_config.add(
+        if transform_container.add(
             EPShardingInfo.from_node(
                 node,
                 rank=rank,
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
index 9cdb20422fb5..5e71cd66c6be 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
+++ b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
@@ -276,6 +276,37 @@ def is_any_lin_op(node: Node) -> bool:
     return is_linear_op(node) or is_fake_quantized_linear_op(node)
 
 
+def is_any_moe_op(node: Node) -> bool:
+    return is_op(
+        node,
+        ops=[
+            torch.ops.auto_deploy.torch_moe,
+            torch.ops.auto_deploy.torch_quant_fp8_moe,
+            torch.ops.auto_deploy.torch_quant_nvfp4_moe,
+            torch.ops.auto_deploy.triton_mxfp4_moe,
+        ],
+    )
+
+
+def is_any_ssm_op(node: Node) -> bool:
+    return is_op(
+        node,
+        ops=[
+            torch.ops.auto_deploy.torch_ssm,
+        ],
+    )
+
+
+def is_any_attention_op(node: Node) -> bool:
+    return is_op(
+        node,
+        ops=[
+            torch.ops.auto_deploy.torch_attention_sdpa,
+            torch.ops.auto_deploy.torch_attention,
+        ],
+    )
+
+
 def is_linear_op(node: Node) -> bool:
     """Check if the node is a linear op.
 
@@ -393,24 +424,54 @@ def identify_regions_between_residuals(gm: GraphModule) -> List[Node]:
 
 
 def get_all_layer_subgraphs(gm: GraphModule) -> List[List[Node]]:
-    """Get subgraphs corresponding to all consecutive layers (attention, MLP, SSM, MoE) in the graph."""
+    """
+    Get subgraphs corresponding to all consecutive layers (attention, MLP, SSM, MoE) in the graph.
+
+    Assumptions:
+        1. each layer (each subgraph) is contained between a list of opening
+        linear layers (e.g., q/k/v_proj, gate/up_proj, in_proj, etc.) and a single closing linear layer
+        (e.g., out_proj, down_proj, etc.)
+        2. all layers are connected in sequence, there are no "parallel" layers.
+
+    Consequence:
+        1. We can linearize both all linear nodes in the graph and the layers itself, having a single
+        linear history and define layers by the indices of corresponding opening and closing linear nodes,
+        with no overlap between layers.
+        E.g., if layer i is defined between linear nodes start_i and end_i, then all linear nodes between
+        start_i and end_i necessarily belong to layer i.
+        2. It does not mean that every linear node in the graph belongs to a layer, that is. Then, if:
+            end_i < start_{{i+1}}, then all linear nodes in[end_i + 1, start_{{i+1}} - 1] do not belong
+            to any layer, and are marked as "unprocessed".
+    Note:
+        The interesting case is MoE with shared experts. In this case, there are two parallel paths:
+        1. Routed experts
+        2. Shared experts
+        In this case, "shared experts" path will be marked as an MLP layer, and "routed experts" path will
+        be "unprocessed". This is desired, since routed experts should not be sharded by the TP transform,
+        but a corresponding EP/BMM transforms.
+    """
 
     assert gm.graph.nodes, "Graph is empty"
     layer_subgraphs = []
-    linear_nodes = list(filtered_nodes(gm.graph.nodes, is_linear_op))
+    linear_nodes = list(filtered_nodes(gm.graph.nodes, is_any_lin_op))
     unprocessed_linear_nodes = set(linear_nodes)
     assert len(linear_nodes) > 0, "Could not find any linear nodes in the graph"
 
     terminating_indices = [-1]
+    last_lin_index = terminating_indices[-1] + 1
 
     # for each linear node, find its layer subgraph defined as regions between consecutive linear nodes
-    while True:
-        layer_subgraph = get_layer_after_linear_node(linear_nodes, terminating_indices)
-        unprocessed_linear_nodes -= set(layer_subgraph)
-        layer_subgraphs.append(layer_subgraph)
+    while last_lin_index < len(linear_nodes):
+        # opening is the list of linear nodes
+        # layer_subgraph is the list of nodes between the opening and closing linear nodes
+        # closing is the last linear node in the layer
+        opening, layer_subgraph, closing = get_layer_after_linear_node(
+            linear_nodes, terminating_indices
+        )
+        if opening is not None:
+            unprocessed_linear_nodes -= set(opening) | set([closing])
+            layer_subgraphs.append([opening, layer_subgraph, closing])
         last_lin_index = terminating_indices[-1] + 1
-        if last_lin_index >= len(linear_nodes):
-            break
 
     # unprocessed linear nodes can be "simple sharded".
     return layer_subgraphs, unprocessed_linear_nodes
@@ -653,22 +714,44 @@ def boundary_condition(n):
 def get_layer_after_linear_node(
     linear_nodes: List[Node], terminating_indices: List[int]
 ) -> List[Node]:
-    """Get the layer after a linear node."""
+    """
+    Get the next model layer.
+    The previous layer was closed by the terminating linear node with index terminating_indices[-1].
+
+    Since we assume a layer is always terminated by a single linear node, we iteratively query subgraph
+    and check for the condition len(lin_nodes_in_subgraph) == 1. If a given linear node
+    linear_nodes[start_lin_index] does not have a corresponding single sink linear node, it will
+    be classified as "unprocessed", and the next linear node is picked as a candidate to open
+    a new layer.
+    """
+
+    def boundary_condition(node: Node) -> bool:
+        return (
+            is_any_lin_op(node)
+            or is_any_moe_op(node)
+            or is_op(node, ops=[torch.ops.aten.sym_size, torch.ops.aten.bmm])
+        )
+
+    def filter_condition(node: Node) -> bool:
+        return is_any_lin_op(node)
+
     lin_nodes_in_subgraph = []
     start_lin_index = terminating_indices[-1] + 1
     while len(lin_nodes_in_subgraph) != 1:
+        if start_lin_index >= len(linear_nodes):
+            terminating_indices.append(len(linear_nodes))
+            return None, None, None
         forward_subgraph = subgraph(
-            sources=[linear_nodes[start_lin_index]], boundary_condition=is_linear_op
-        )
-        lin_nodes_in_subgraph = list(
-            filtered_nodes(forward_subgraph, ops=torch.ops.auto_deploy.torch_linear_simple)
+            sources=[linear_nodes[start_lin_index]], boundary_condition=boundary_condition
         )
+        lin_nodes_in_subgraph = list(filtered_nodes(forward_subgraph, filter_condition))
         start_lin_index += 1
     terminating_linear_node = lin_nodes_in_subgraph[0]
-    # get all opening linear nodes
-    opening_linear_nodes = subgraph(
-        sinks=[terminating_linear_node], include=is_linear_op, boundary_condition=is_linear_op
+    backward_subgraph = subgraph(
+        sinks=[terminating_linear_node], boundary_condition=boundary_condition
     )
+    # get all opening linear nodes
+    opening_linear_nodes = list(filtered_nodes(backward_subgraph, is_any_lin_op))
     # opening nodes must succeed last terminating node
     last_terminating_index = terminating_indices[-1]
     opening_linear_nodes = [
@@ -677,7 +760,7 @@ def get_layer_after_linear_node(
     assert linear_nodes[start_lin_index - 1] in opening_linear_nodes, (
         "Linear node not found in opening linear nodes"
     )
-    layer_subgraph = opening_linear_nodes + forward_subgraph
+    layer_subgraph = [opening_linear_nodes, backward_subgraph, terminating_linear_node]
 
     # return the index of the terminating linear node
     if terminating_linear_node == linear_nodes[-1]:
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py
index d12b99545440..bd576cfeba74 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py
+++ b/tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py
@@ -163,21 +163,22 @@ def _validate_sharded_shapes(
         include=lambda n: is_op(n, [torch.ops.aten.view, torch.ops.aten.reshape]),
         boundary_condition=is_any_lin_op,
     )
-    for view_node in nodes_to_validate:
-        if len(view_node.args) < 2:
+    for shape_node in nodes_to_validate:
+        # Parameter update must be idempotent
+        if "sharded" in shape_node.meta and shape_node.meta["sharded"]:
             continue
-        if "sharded" in view_node.meta and view_node.meta["sharded"]:
+        if len(shape_node.args) < 2:
             continue
-        view_shape = list(view_node.args[1])
+        view_shape = list(shape_node.args[1])
         if not isinstance(view_shape, list):
             continue
         if len(view_shape) >= 3 and isinstance(view_shape[2], int) and view_shape[2] != -1:
-            args = list(view_node.args)
+            args = list(shape_node.args)
             view_shape[2] = -1  # view_shape[2] // world_size
             args[1] = tuple(view_shape)
-            view_node.args = tuple(args)
-            view_node.meta["sharded"] = True
-            ad_logger.debug(f"\nUpdated view node {view_node} arguments to {view_node.args}")
+            shape_node.args = tuple(args)
+            shape_node.meta["sharded"] = True
+            ad_logger.debug(f"\nUpdated view node {shape_node} arguments to {shape_node.args}")
 
     # if fused_weight_dims is provided, we need to update all split sizes
     if fused_weight_dims is not None:
@@ -187,14 +188,18 @@ def _validate_sharded_shapes(
         split_nodes = subgraph(
             [node],
             [next_lin_node],
-            include=lambda n: is_op(n, [torch.ops.aten.split, torch.ops.aten.split_with_sizes]),
+            include=lambda n: is_op(n, [torch.ops.aten.split_with_sizes]),
         )
         for split_node in split_nodes:
+            # Parameter update must be idempotent
+            if "sharded" in split_node.meta and split_node.meta["sharded"]:
+                continue
             orig_sizes = split_node.args[1]
             new_sizes = [orig_sizes[i] // world_size for i in range(len(orig_sizes))]
             args = list(split_node.args)
             args[1] = new_sizes
             split_node.args = tuple(args)
+            split_node.meta["sharded"] = True
             ad_logger.debug(f"\nUpdated split node {split_node} arguments to {split_node.args}")
 
 
@@ -256,10 +261,6 @@ def split_fused_tensor(
             fused_dims: list = fused_weight_dims,
             d: int = dim,
         ) -> torch.Tensor:
-            # dim_d = t.shape[d]
-            # num_parts = 1
-            # part_size = dim_d // num_parts
-            # fused_dims = [part_size] * num_parts
             return torch.cat(
                 [split_tensor(w) for w in torch.split(t, fused_dims, dim=d)],
                 dim=d,
@@ -407,7 +408,9 @@ def _insert_sharded_mamba(
     conv_args[-1] = conv1d_node.args[-1] // world_size
     conv1d_node.args = tuple(conv_args)
 
-    # First, shard the entry_node (the first linear layer)
+    ##############################################################
+    ####### shard the entry_node (the first linear layer) ########
+    ##############################################################
     # Extract entry node's fused_weight_dims by matching weight name against patterns
     entry_fused_dims = None
     if fused_weight_dims:
@@ -431,6 +434,9 @@ def _insert_sharded_mamba(
         allreduce_strategy=allreduce_strategy,
     )
 
+    ##############################################################
+    ######## Shard remaining weights: conv1d and RMSNorm #########
+    ##############################################################
     # Get all weight nodes in the subgraph except for out_proj
     weight_nodes = [
         n
@@ -438,7 +444,6 @@ def _insert_sharded_mamba(
         if "out_proj" not in str(n)
     ]
 
-    # Shard remaining weights, such as conv1d or RMSNorm
     for weight_node in weight_nodes:
         weight_key = weight_node.target
 
@@ -481,6 +486,45 @@ def _insert_sharded_mamba(
             f"{weight_param.shape} -> {sharded_shape}"
         )
 
+    ##############################################################
+    ############## update split node parameters ##################
+    ##############################################################
+    next_lin_node, _ = bfs(entry_node, is_any_lin_op, include_root=False)
+
+    split_nodes = subgraph(
+        [entry_node],
+        [next_lin_node],
+        include=lambda n: is_op(n, [torch.ops.aten.split_with_sizes]),
+    )
+    for split_node in split_nodes:
+        orig_sizes = split_node.args[1]
+        new_sizes = [orig_sizes[i] // world_size for i in range(len(orig_sizes))]
+        args = list(split_node.args)
+        args[1] = new_sizes
+        split_node.args = tuple(args)
+        ad_logger.debug(f"\nUpdated split node {split_node} arguments to {split_node.args}")
+
+    nodes_to_validate = subgraph(
+        [entry_node],
+        include=lambda n: is_op(n, [torch.ops.aten.view, torch.ops.aten.reshape]),
+        boundary_condition=is_any_lin_op,
+    )
+    for reshape_node in nodes_to_validate:
+        if len(reshape_node.args) < 2:
+            continue
+        if "sharded" in reshape_node.meta and reshape_node.meta["sharded"]:
+            continue
+        view_shape = list(reshape_node.args[1])
+        if not isinstance(view_shape, list):
+            continue
+        if len(view_shape) >= 3 and isinstance(view_shape[2], int) and view_shape[2] != -1:
+            args = list(reshape_node.args)
+            view_shape[2] = -1  # view_shape[2] // world_size
+            args[1] = tuple(view_shape)
+            reshape_node.args = tuple(args)
+            reshape_node.meta["sharded"] = True
+            ad_logger.debug(f"\nUpdated view node {reshape_node} arguments to {reshape_node.args}")
+
 
 def _shard_parameter_node(
     gm: GraphModule,
@@ -574,10 +618,6 @@ def _shard_parameter_node(
 
     # # # column shard with no gather: the output is sharded
     if not add_dist:
-        if is_any_lin_op(node):
-            _validate_sharded_shapes(
-                node, fused_weight_dims=fused_weight_dims, world_size=world_size
-            )
         return
 
     # figure out the right dist op (backend-aware)
@@ -717,36 +757,19 @@ def validate(self, gm: GraphModule = None, node: Node = None) -> bool:
 
     def apply(self, gm: GraphModule, node: Node) -> None:
         """Apply TP sharding transformation to the graph module."""
-        if self.layer_type == LayerType.MAMBA:
-            _insert_sharded_mamba(
-                gm=gm,
-                entry_node=node,
-                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
-                if isinstance(self.fused_weight_dims, dict)
-                else None,
-                quantization_cb=self.quantization_cb,
-                allreduce_strategy=self.allreduce_strategy,
-            )
-        else:
-            _shard_parameter_node(
-                gm=gm,
-                node=node,
-                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,
-                quantization_cb=self.quantization_cb,
-                allreduce_strategy=self.allreduce_strategy,
-            )
+        _shard_parameter_node(
+            gm=gm,
+            node=node,
+            dim=self.split_dim.value,
+            rank=self.rank,
+            world_size=self.world_size,
+            add_dist=self.dist_op is not None,
+            dist_backend=self.dist_backend,
+            min_local_shape=self.min_local_shape,
+            fused_weight_dims=self.fused_weight_dims,
+            quantization_cb=self.quantization_cb,
+            allreduce_strategy=self.allreduce_strategy,
+        )
 
 
 class ParameterUpdateInfo(ShardingTransformInfo):
@@ -1337,7 +1360,6 @@ class ShardingSource(Enum):
 class ShardingDim(Enum):
     """Enum for sharding dimension."""
 
-    SSM = "ssm"
     TP = "tp"
     EP = "ep"
     BMM = "bmm"
@@ -1365,7 +1387,7 @@ class ShardingTransformContainer(BaseModel):
         default_factory=lambda: [ShardingSource.HEURISTIC]
     )
     sharding_dims: List[ShardingDim] = Field(
-        default_factory=lambda: [ShardingDim.SSM, ShardingDim.TP, ShardingDim.EP, ShardingDim.BMM]
+        default_factory=lambda: [ShardingDim.TP, ShardingDim.EP, ShardingDim.BMM]
     )
     allreduce_strategy: AllReduceStrategy = Field(
         default=AllReduceStrategy.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 d34c8c0f3d20..9a8f0d5164ee 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
@@ -19,7 +19,7 @@
 )
 from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer
 from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_linear_op, is_op
-from tensorrt_llm._torch.auto_deploy.utils.sharding_utils import FP8TPShardingInfo
+from tensorrt_llm._torch.auto_deploy.utils.sharding_utils import FP8TPShardingInfo, LayerType
 from tensorrt_llm.functional import AllReduceStrategy
 
 base_model_tp_plan = {
@@ -123,7 +123,7 @@ def forward(self, x):
         return self.linear2(y)
 
 
-def _run_job(
+def _run_sharding_execution_job(
     model_cls: nn.Module,
     dist_op_expected: str,
     bias: bool,
@@ -212,7 +212,6 @@ def combined_graph_check(gm) -> bool:
         has_expected_dist_ops = any(is_op(n, op_expected) for n in gm.graph.nodes) == (
             world_size > 1
         )
-        # Check weight size constraints
         weight_sizes_valid = verify_local_weight_sizes(gm)
         return has_expected_dist_ops and weight_sizes_valid
 
@@ -282,6 +281,7 @@ def _run_pattern_detection_job(
                             min_local_shape=min_local_shape,
                             allreduce_strategy=AllReduceStrategy.AUTO,
                             dist_backend="auto",
+                            layer_type=LayerType.ATTENTION,
                         )
                     )
         elif model_cls == MLP:
@@ -390,7 +390,7 @@ def test_sharding(
     from_config: bool,
 ):
     dist_common.spawn_multiprocess_job(
-        job=partial(_run_job, model_cls, dist_op_expected, bias, from_config),
+        job=partial(_run_sharding_execution_job, model_cls, dist_op_expected, bias, from_config),
         size=device_count,
     )
 
@@ -420,7 +420,3 @@ def test_sharding_pattern_detection(
     No need to run distributed job, can be run on single process.
     """
     _run_pattern_detection_job(model_cls, bias, 0, world_size, from_config)
-
-
-if __name__ == "__main__":
-    _run_pattern_detection_job(nn.Linear, False, 0, 8, False)