From ec2927bc64f96a61eab257e2fd8c80229b581982 Mon Sep 17 00:00:00 2001
From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Date: Wed, 17 Dec 2025 05:47:43 -0800
Subject: [PATCH 1/3] By default, unmatched lin nodes are not sharded

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
---
 tensorrt_llm/_torch/auto_deploy/config/default.yaml           | 1 +
 tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py | 4 +++-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml
index 4200be6c6aab..38d271ffacc3 100644
--- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml
+++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml
@@ -81,6 +81,7 @@ transforms:
     sharding_source: ['manual', 'factory', 'heuristic']
     support_partial_config: true
     sharding_dims: ['tp', 'ep', 'bmm']
+    shard_all_unprocessed: false
     allreduce_strategy: 'NCCL'
     dist_backend: auto
     requires_shape_prop: true
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
index bae85f3a2216..00272f6a3c3a 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
@@ -140,6 +140,7 @@ class ShardingTransformConfig(TransformConfig):
     sharding_dims: List[ShardingDim] = Field(
         default_factory=lambda: [ShardingDim.TP, ShardingDim.EP, ShardingDim.BMM]
     )
+    shard_all_unprocessed: bool = Field(default=False)
     allreduce_strategy: AllReduceStrategy = Field(
         default=AllReduceStrategy.AUTO,
         description="AllReduce strategy for distributed operations. "
@@ -2608,7 +2609,8 @@ def detect_column_row_shard(
                 num_attention_shards += 1
 
     # simple shard remaining linear nodes
-    num_simple_shards += _process_simple_shard(unprocessed_linear_nodes, transform_container)
+    if config.shard_all_unprocessed:
+        num_simple_shards += _process_simple_shard(unprocessed_linear_nodes, transform_container)
     num_column_row_shards += num_ssm_shards
     num_shards = num_simple_shards + num_column_row_shards
     ad_logger.info(

From 3c04dd91443378e066e0bed1bca0934d3b6bc5ae Mon Sep 17 00:00:00 2001
From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Date: Wed, 17 Dec 2025 06:18:55 -0800
Subject: [PATCH 2/3] updated documentatino

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
---
 .../_torch/auto_deploy/transform/library/sharding.py        | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
index 00272f6a3c3a..044cf7002079 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
@@ -140,7 +140,11 @@ class ShardingTransformConfig(TransformConfig):
     sharding_dims: List[ShardingDim] = Field(
         default_factory=lambda: [ShardingDim.TP, ShardingDim.EP, ShardingDim.BMM]
     )
-    shard_all_unprocessed: bool = Field(default=False)
+    shard_all_unprocessed: bool = Field(
+        default=False,
+        description="When True, apply simple shard (column split + all_gather) to "
+        "'leftover' linear nodes that are not part of any layer subgraph.",
+    )
     allreduce_strategy: AllReduceStrategy = Field(
         default=AllReduceStrategy.AUTO,
         description="AllReduce strategy for distributed operations. "

From 3a96a61943355a836d5dc1977545b2f275b26ff3 Mon Sep 17 00:00:00 2001
From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Date: Mon, 22 Dec 2025 07:49:43 -0800
Subject: [PATCH 3/3] Restored default behavior

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
---
 tensorrt_llm/_torch/auto_deploy/config/default.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml
index 38d271ffacc3..bdc97520fac3 100644
--- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml
+++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml
@@ -81,7 +81,7 @@ transforms:
     sharding_source: ['manual', 'factory', 'heuristic']
     support_partial_config: true
     sharding_dims: ['tp', 'ep', 'bmm']
-    shard_all_unprocessed: false
+    shard_all_unprocessed: true
     allreduce_strategy: 'NCCL'
     dist_backend: auto
     requires_shape_prop: true