From d5025f25f2e65edabc8aad3aa2ea7d2c8244f4dc Mon Sep 17 00:00:00 2001
From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Date: Fri, 20 Feb 2026 04:19:22 -0800
Subject: [PATCH 1/7] added GatedDeltaNet support from config

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
---
 .../configs/qwen3.5_moe_35b.yaml              | 20 +++++++
 .../configs/qwen3.5_moe_400b.yaml             | 19 ++++++
 .../auto_deploy/transform/library/sharding.py | 16 ++++-
 .../_torch/auto_deploy/utils/node_utils.py    |  4 +-
 .../library/test_tp_sharding.py               | 59 +++++++++----------
 5 files changed, 82 insertions(+), 36 deletions(-)

diff --git a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml
index bf1a5593e144..1a86f46262c6 100644
--- a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml
+++ b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml
@@ -16,3 +16,23 @@ model_kwargs:
   #   num_hidden_layers: 6
   # vision_config:
   #   depth: 2
+transforms:
+  detect_sharding:
+    sharding_dims: ['tp','ep', 'bmm']
+    # use only manual config for TP sharding
+    sharding_source: ['manual']
+    manual_config:
+      tp_plan:
+        # GDN layer
+        "in_proj_qkv": "delta"
+        # attention layer
+        "q_proj": "colwise"
+        "k_proj": "colwise"
+        "v_proj": "colwise"
+        "o_proj": "rowwise"
+        # replicating shared experts (keep them commented out)
+        # "shared_expert_gate_proj": "colwise"
+        # "shared_expert_up_proj": "colwise"
+        # "shared_expert_down_proj": "rowwise"
+        # gating layer should be replicated as well
+        # "gate": "gather"
diff --git a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml
index d6462d1ca52d..12175d03d673 100644
--- a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml
+++ b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml
@@ -16,3 +16,22 @@ model_kwargs:
 transforms:
   export_to_gm:
     num_moe_experts_for_export: 2
+  detect_sharding:
+    sharding_dims: ['tp','ep', 'bmm']
+    # use only manual config for TP sharding
+    sharding_source: ['manual']
+    manual_config:
+      tp_plan:
+        # GDN layer
+        "in_proj_qkv": "delta"
+        # attention layer
+        "q_proj": "colwise"
+        "k_proj": "colwise"
+        "v_proj": "colwise"
+        "o_proj": "rowwise"
+        # replicating shared experts (keep them commented out)
+        # "shared_expert_gate_proj": "colwise"
+        # "shared_expert_up_proj": "colwise"
+        # "shared_expert_down_proj": "rowwise"
+        # gating layer should be replicated as well
+        # "gate": "gather"
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
index 3c9398a946e3..8bc5d2ad6edd 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
@@ -239,7 +239,9 @@ def validate_config(self, sources: Union[ShardingSource, List[ShardingSource]] =
             supported_modes = {
                 "colwise",  # row split and no collective
                 "rowwise",  # column split and all-reduce
+                "mla",  # mla layer
                 "mamba",  # mamba SSM layer
+                "delta",  # gated delta net layer
                 "gather",  # simple shard (row + all_gather)
                 # TODO: remaining values are not supported yet.
                 # They require hybrid EP+TP and/or SP support.
@@ -2663,6 +2665,8 @@ def detect_sharding_from_config(
     num_row_col_shards = 0
     num_attention_shards = 0
     num_ssm_shards = 0
+    num_mla_shards = 0
+    num_delta_shards = 0
     linear_nodes = list(filtered_nodes(gm.graph.nodes, is_any_lin_op))
 
     # use layer_subgraphs to determine the layer_type
@@ -2724,7 +2728,14 @@ def detect_sharding_from_config(
                     if _process_ssm_sharding(layer_subgraph, transform_container) > 0:
                         num_ssm_shards += 1
                         num_row_col_shards += 1
-
+                elif config == "gdn":
+                    if _process_delta_sharding(layer_subgraph, transform_container) > 0:
+                        num_delta_shards += 1
+                        num_row_col_shards += 1
+                elif config == "mla":
+                    if _process_mla_sharding(layer_subgraph, transform_container) > 0:
+                        num_mla_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.")
@@ -2783,7 +2794,8 @@ def detect_sharding_from_config(
     num_shards = num_simple_shards + num_row_col_shards
     ad_logger.info(
         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})"
+        f"row-col: {num_row_col_shards} (attention: {num_attention_shards}, "
+        f"mla: {num_mla_shards}, delta: {num_delta_shards}, ssm: {num_ssm_shards})"
     )
 
     num_matches = len(transform_container.weight_sharding_transforms)
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
index ddd38e4e9629..22ef6c26be2f 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
+++ b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
@@ -45,11 +45,11 @@ class LayerType(Enum):
 
 class LayerSubgraph(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
+    layer_type: LayerType
     opening_nodes: List[Node]
-    subgraph_nodes: List[Node]
     terminating_node: Union[Node, None]
-    layer_type: LayerType
     min_local_shape: int = 1
+    subgraph_nodes: List[Node]
 
 
 class WeightNode(BaseModel):
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 f42d4856800c..1dd33fd81c74 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
@@ -27,36 +27,26 @@
 from tensorrt_llm.functional import AllReduceStrategy
 
 base_model_tp_plan = {
-    "q_proj": "colwise",
-    "k_proj": "colwise",
-    "v_proj": "colwise",
-    "o_proj": "rowwise",
-    "gate_proj": "colwise",
-    "up_proj": "colwise",
-    "down_proj": "rowwise",
-    "linear1": "colwise",
-    "linear2": "rowwise",
-    "linear": "gather",
-    # Mamba2 specific projections
-    "in_proj": "mamba",
-    "out_proj": "rowwise",
-    # MLA specific projections
-    "q_a_proj": "gather",
-    "q_b_proj": "colwise",
-    "kv_a_proj_with_mqa": "gather",
-    "kv_b_proj": "colwise",
-    # "input_layernorm.weight": "sequence_parallel",
-    # "post_attention_layernorm.weight": "sequence_parallel",
-    # "norm.weight": "sequence_parallel",
-    # "shared_expert.gate_proj": "local_colwise",
-    # "shared_expert.up_proj": "local_colwise",
-    # "shared_expert.down_proj": "local_rowwise",
-    # "experts.gate_up_proj": "local_packed_rowwise",
-    # "experts.down_proj": "local_colwise",
-    # "experts": "local",
-    "feed_forward": "gather",
-    "self": "gather",
-    "weight": "gather",
+    # ATTENTION
+    # "q_proj": "colwise",
+    # "k_proj": "colwise",
+    # "v_proj": "colwise",
+    # "o_proj": "rowwise",
+    # # gated MLP
+    # "gate_proj": "colwise",
+    # "up_proj": "colwise",
+    # "down_proj": "rowwise",
+    # # GPT-style MLP
+    # "linear1": "colwise",
+    # "linear2": "rowwise",
+    # "linear": "gather",
+    # # Mamba2 specific projections
+    # "in_proj": "mamba",
+    # # MLA specific projections
+    # "q_a_proj": "mla",
+    # # delta net
+    # # qkv regex matches both fused (qkvz) and unfused (qkv) opening nodes
+    # "in_proj_qkv": "delta",
 }
 
 predefined_config = {
@@ -563,13 +553,16 @@ def verify_local_weight_sizes(gm) -> bool:
     op_expected = getattr(torch.ops.auto_deploy, dist_op_expected)
 
     gm = torch_export_to_gm(model, args=(x,), clone=True)
+
+    sharding_source = "heuristic" if from_config else "manual"
     gm_transformed = InferenceOptimizer(
         None,
         {
             "detect_sharding": {
                 "stage": "sharding",
                 "sharding_scope": ["tp"],
-                "tp_sharding_source": ["heuristic"],
+                "sharding_source": [sharding_source],
+                "manual_config": predefined_config,
             },
             "sharding_transform_executor": {
                 "stage": "sharding",
@@ -931,6 +924,7 @@ def _run_pattern_detection_job(
                             )
                         )
 
+    sharding_source = "heuristic" if from_config else "manual"
     # get detected transformations
     optimizer = InferenceOptimizer(
         None,
@@ -938,8 +932,9 @@ def _run_pattern_detection_job(
             "detect_sharding": {
                 "stage": "sharding",
                 "sharding_scope": ["tp"],
-                "tp_sharding_source": ["heuristic"],
+                "sharding_source": [sharding_source],
                 "shard_all_unprocessed": True,
+                "manual_config": predefined_config,
             },
         },
     )

From 4812b3f9d1e37f3282b04bfc40bcabaf7ebf1f76 Mon Sep 17 00:00:00 2001
From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Date: Fri, 20 Feb 2026 04:47:34 -0800
Subject: [PATCH 2/7] fixed config keyword

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
---
 tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py | 2 +-
 1 file changed, 1 insertion(+), 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 8bc5d2ad6edd..bc56a877a608 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
@@ -2728,7 +2728,7 @@ def detect_sharding_from_config(
                     if _process_ssm_sharding(layer_subgraph, transform_container) > 0:
                         num_ssm_shards += 1
                         num_row_col_shards += 1
-                elif config == "gdn":
+                elif config == "delta":
                     if _process_delta_sharding(layer_subgraph, transform_container) > 0:
                         num_delta_shards += 1
                         num_row_col_shards += 1

From cd81fcae0fbb53330695eefb5669eb5afc9cec42 Mon Sep 17 00:00:00 2001
From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Date: Fri, 20 Feb 2026 04:52:23 -0800
Subject: [PATCH 3/7] updated test configs

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
---
 .../library/test_tp_sharding.py               | 38 +++++++++----------
 1 file changed, 19 insertions(+), 19 deletions(-)

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 1dd33fd81c74..37ff829e837a 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
@@ -28,25 +28,25 @@
 
 base_model_tp_plan = {
     # ATTENTION
-    # "q_proj": "colwise",
-    # "k_proj": "colwise",
-    # "v_proj": "colwise",
-    # "o_proj": "rowwise",
-    # # gated MLP
-    # "gate_proj": "colwise",
-    # "up_proj": "colwise",
-    # "down_proj": "rowwise",
-    # # GPT-style MLP
-    # "linear1": "colwise",
-    # "linear2": "rowwise",
-    # "linear": "gather",
-    # # Mamba2 specific projections
-    # "in_proj": "mamba",
-    # # MLA specific projections
-    # "q_a_proj": "mla",
-    # # delta net
-    # # qkv regex matches both fused (qkvz) and unfused (qkv) opening nodes
-    # "in_proj_qkv": "delta",
+    "q_proj": "colwise",
+    "k_proj": "colwise",
+    "v_proj": "colwise",
+    "o_proj": "rowwise",
+    # gated MLP
+    "gate_proj": "colwise",
+    "up_proj": "colwise",
+    "down_proj": "rowwise",
+    # GPT-style MLP
+    "linear1": "colwise",
+    "linear2": "rowwise",
+    "linear": "gather",
+    # Mamba2 specific projections
+    "in_proj": "mamba",
+    # MLA specific projections
+    "q_a_proj": "mla",
+    # delta net
+    # qkv regex matches both fused (qkvz) and unfused (qkv) opening nodes
+    "in_proj_qkv": "delta",
 }
 
 predefined_config = {

From 62062b50407bbe4cb88bababc0b0cdbf1b2ec551 Mon Sep 17 00:00:00 2001
From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Date: Fri, 20 Feb 2026 05:01:44 -0800
Subject: [PATCH 4/7] up

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
---
 .../_torch/auto_deploy/distributed/common.py  | 187 +++++++++---------
 1 file changed, 92 insertions(+), 95 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/distributed/common.py b/tensorrt_llm/_torch/auto_deploy/distributed/common.py
index 9f1dc85f1063..4a92c6ee4f38 100644
--- a/tensorrt_llm/_torch/auto_deploy/distributed/common.py
+++ b/tensorrt_llm/_torch/auto_deploy/distributed/common.py
@@ -2,6 +2,7 @@
 
 import atexit
 import os
+import socket
 import sys
 from typing import Callable, List, Optional, Tuple
 
@@ -85,16 +86,16 @@ def initialize_or_skip(
     rank: int = 0,
     world_size: int = 1,
     port: Optional[int] = None,
-    shared_port: Optional["mp.Value"] = None,
-    port_ready_barrier: Optional["mp.Barrier"] = None,
+    port_recv_conn=None,
+    port_send_conns=None,
 ) -> Tuple[int, int]:
     if not dist.is_initialized():
         return initialize(
             rank=rank,
             world_size=world_size,
             port=port,
-            shared_port=shared_port,
-            port_ready_barrier=port_ready_barrier,
+            port_recv_conn=port_recv_conn,
+            port_send_conns=port_send_conns,
         )
     return get_rank(), get_world_size()
 
@@ -129,32 +130,23 @@ def _set_distributed_env_vars(local_rank: int, world_size: int, port: int) -> No
     os.environ["LOCAL_RANK"] = str(local_rank)
 
 
-def _try_init_process_group(local_rank: int, world_size: int, port: int) -> bool:
-    """Attempt to initialize process group. Returns True on success, False on EADDRINUSE."""
-    _set_distributed_env_vars(local_rank, world_size, port)
-
+def _is_port_available(port: int) -> bool:
+    """Lightweight check: try to bind to the port and release immediately."""
     try:
-        dist.init_process_group(
-            "nccl",
-            world_size=world_size,
-            rank=local_rank,
-            device_id=torch.device(local_rank),
-        )
-        return True
-    except Exception as e:
-        # Check if this is a port-in-use error (only rank 0 binds, so only rank 0 can get this)
-        if "EADDRINUSE" in str(e) or "address already in use" in str(e).lower():
-            ad_logger.warning(f"Port {port} already in use, will retry with new port")
-            return False
-        raise
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            s.bind(("127.0.0.1", port))
+            return True
+    except OSError:
+        return False
 
 
 def initialize(
     rank: int = 0,
     world_size: int = 1,
     port: Optional[int] = None,
-    shared_port: Optional["mp.Value"] = None,
-    port_ready_barrier: Optional["mp.Barrier"] = None,
+    port_recv_conn=None,
+    port_send_conns=None,
     max_retries: int = 5,
 ) -> Tuple[int, int]:
     """Initialize distributed process group.
@@ -163,8 +155,8 @@ def initialize(
         rank: Process rank (ignored for OMPI/torchelastic).
         world_size: Total number of processes (ignored for OMPI/torchelastic).
         port: Initial port to try. If None, a free port will be selected.
-        shared_port: Optional mp.Value for rank 0 to share the final port with other ranks.
-        port_ready_barrier: Optional mp.Barrier to synchronize port selection.
+        port_recv_conn: Pipe read-end for this rank to receive the resolved port from rank 0.
+        port_send_conns: List of Pipe write-ends for rank 0 to broadcast the port.
         max_retries: Maximum number of port retry attempts for rank 0.
     """
     if is_ompi():
@@ -189,65 +181,57 @@ def initialize(
     # Necessary to assign a device to each rank.
     torch.cuda.set_device(local_rank)
 
-    # If we have shared port synchronization (multiprocess spawn mode)
-    if shared_port is not None and port_ready_barrier is not None:
-        if local_rank == 0:
-            # Rank 0: try ports until one works, then share with other ranks
-            init_success = False
-            init_error = None
-            try:
-                for attempt in range(max_retries):
-                    ad_logger.info(
-                        f"Initializing for: {lib=}, {local_rank=}, {world_size=}, {port=} (attempt {attempt + 1})"
-                    )
-                    if _try_init_process_group(local_rank, world_size, port):
-                        # Success! Share the working port with other ranks
-                        shared_port.value = port
-                        init_success = True
-                        break
-                    else:
-                        # Port was taken, try a new one
-                        port = get_free_port()
-                else:
-                    # All retries exhausted
-                    init_error = RuntimeError(
-                        f"Failed to find available port after {max_retries} attempts"
-                    )
-            except Exception as e:
-                # Catch any unexpected error so we can still signal other ranks
-                init_error = e
-            finally:
-                # ALWAYS signal other ranks, even on error, to prevent deadlock
-                if not init_success:
-                    shared_port.value = -1
-                port_ready_barrier.wait()
-
-            if init_error is not None:
-                raise init_error
-        else:
-            # Other ranks: wait for rank 0 to find a working port
-            port_ready_barrier.wait()
-            port = shared_port.value
-            if port == -1:
-                raise RuntimeError("Rank 0 failed to initialize, cannot proceed")
-            ad_logger.info(f"Initializing for: {lib=}, {local_rank=}, {world_size=}, {port=}")
-            _set_distributed_env_vars(local_rank, world_size, port)
-            dist.init_process_group(
-                "nccl",
-                world_size=world_size,
-                rank=local_rank,
-                device_id=torch.device(local_rank),
-            )
-    else:
-        # Original path: no retry mechanism (OMPI, torchelastic, or single process)
-        ad_logger.info(f"Initializing for: {lib=}, {local_rank=}, {world_size=}, {port=}")
-        _set_distributed_env_vars(local_rank, world_size, port)
-        dist.init_process_group(
-            "nccl",
-            world_size=world_size,
-            rank=local_rank,
-            device_id=torch.device(local_rank),
-        )
+    # Pipe-based port synchronization (multiprocess spawn mode).
+    # Port resolution MUST happen before init_process_group because that call
+    # is collective — all ranks must participate simultaneously.  Rank 0 does a
+    # lightweight socket-bind check (not a full init) to find a usable port,
+    # broadcasts it through the pipes, and then ALL ranks call
+    # init_process_group together.
+    if port_send_conns is not None:
+        assert local_rank == 0, "port_send_conns should only be provided to rank 0"
+        port_ok = False
+        init_error = None
+        try:
+            for attempt in range(max_retries):
+                ad_logger.info(
+                    f"Checking port for: {lib=}, {local_rank=}, {world_size=}, {port=} (attempt {attempt + 1})"
+                )
+                if _is_port_available(port):
+                    port_ok = True
+                    break
+                ad_logger.warning(f"Port {port} already in use, will retry with new port")
+                port = get_free_port()
+            else:
+                init_error = RuntimeError(
+                    f"Failed to find available port after {max_retries} attempts"
+                )
+        except Exception as e:
+            init_error = e
+        finally:
+            # ALWAYS broadcast to other ranks, even on error, to prevent deadlock
+            final_port = port if port_ok else -1
+            for conn in port_send_conns:
+                conn.send(final_port)
+                conn.close()
+
+        if init_error is not None:
+            raise init_error
+    elif port_recv_conn is not None:
+        # Non-rank-0: wait for rank 0 to broadcast the working port
+        port = port_recv_conn.recv()
+        port_recv_conn.close()
+        if port == -1:
+            raise RuntimeError("Rank 0 failed to initialize, cannot proceed")
+
+    # All paths converge here: initialize the process group collectively
+    ad_logger.info(f"Initializing for: {lib=}, {local_rank=}, {world_size=}, {port=}")
+    _set_distributed_env_vars(local_rank, world_size, port)
+    dist.init_process_group(
+        "nccl",
+        world_size=world_size,
+        rank=local_rank,
+        device_id=torch.device(local_rank),
+    )
 
     # Register cleanup function to be called at exit
     atexit.register(cleanup)
@@ -259,11 +243,11 @@ def initialize(
 
 
 def init_and_run_process(
-    job, rank, size, port, shared_port=None, port_ready_barrier=None, **kwargs
+    job, rank, size, port, port_recv_conn=None, port_send_conns=None, **kwargs
 ):
     try:
         initialize_or_skip(
-            rank, size, port, shared_port=shared_port, port_ready_barrier=port_ready_barrier
+            rank, size, port, port_recv_conn=port_recv_conn, port_send_conns=port_send_conns
         )
         job(rank, size, **kwargs)
     except Exception as e:
@@ -318,11 +302,17 @@ def _start_multiprocess_job(
     ctx = mp.get_context("spawn")
     processes: List[mp.Process] = []
 
-    # Create shared state for port synchronization with retry mechanism:
-    # - shared_port: rank 0 writes the final working port here
-    # - port_ready_barrier: all ranks wait here until rank 0 has bound successfully
-    shared_port = ctx.Value("i", port)  # 'i' = signed int
-    port_ready_barrier = ctx.Barrier(size)
+    # Create pipes for port synchronization with retry mechanism.
+    # Pipes use sockets (not POSIX named semaphores like Value/Barrier),
+    # so they work reliably on all filesystems including network mounts.
+    # Rank 0 finds a working port and sends it through each pipe;
+    # other ranks block on recv() until rank 0 is ready.
+    port_recv_conns = []
+    port_send_conns = []
+    for _ in range(size - 1):
+        recv_conn, send_conn = ctx.Pipe(duplex=False)
+        port_recv_conns.append(recv_conn)
+        port_send_conns.append(send_conn)
 
     for rank in range(size):
         if input_queues:
@@ -330,10 +320,15 @@ def _start_multiprocess_job(
         if output_queue:
             kwargs["output_queue"] = output_queue if rank == 0 else None
 
+        rank_kwargs = {
+            **kwargs,
+            "port_recv_conn": port_recv_conns[rank - 1] if rank > 0 else None,
+            "port_send_conns": port_send_conns if rank == 0 else None,
+        }
         p = ctx.Process(
             target=init_and_run_process,
             args=(job, rank, size, port),
-            kwargs={**kwargs, "shared_port": shared_port, "port_ready_barrier": port_ready_barrier},
+            kwargs=rank_kwargs,
             daemon=True,
         )
         p.start()
@@ -350,9 +345,11 @@ def _join_multiprocess_job(processes):
     for p in processes:
         p.join()
 
-        # Ensure that all processes have exited successfully
-        if isinstance(p, mp.Process):
-            assert not p.exitcode, f"Process {p.pid} exited with code {p.exitcode}"
+        # Ensure that all processes have exited successfully.
+        # Check exitcode via hasattr rather than isinstance(p, mp.Process), because
+        # spawn-context processes (SpawnProcess) don't inherit from mp.Process.
+        if hasattr(p, "exitcode"):
+            assert p.exitcode == 0, f"Process {p.pid} exited with code {p.exitcode}"
 
 
 def spawn_multiprocess_job(job: Callable[[int, int], None], size: Optional[int] = None):

From 017260d888e26944210b55e1d9fb39ffb6d3181b Mon Sep 17 00:00:00 2001
From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Date: Sat, 21 Feb 2026 11:39:06 -0800
Subject: [PATCH 5/7] up

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
---
 .../configs/qwen3.5_moe_35b.yaml              |  20 ----
 .../configs/qwen3.5_moe_400b.yaml             |  19 ---
 .../auto_deploy/transform/library/sharding.py |  25 ++--
 .../_torch/auto_deploy/utils/node_utils.py    |  15 ++-
 .../library/test_ep_sharding.py               |  17 ++-
 .../library/test_tp_sharding.py               | 110 +++++++++++++-----
 6 files changed, 110 insertions(+), 96 deletions(-)

diff --git a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml
index 1a86f46262c6..bf1a5593e144 100644
--- a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml
+++ b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_35b.yaml
@@ -16,23 +16,3 @@ model_kwargs:
   #   num_hidden_layers: 6
   # vision_config:
   #   depth: 2
-transforms:
-  detect_sharding:
-    sharding_dims: ['tp','ep', 'bmm']
-    # use only manual config for TP sharding
-    sharding_source: ['manual']
-    manual_config:
-      tp_plan:
-        # GDN layer
-        "in_proj_qkv": "delta"
-        # attention layer
-        "q_proj": "colwise"
-        "k_proj": "colwise"
-        "v_proj": "colwise"
-        "o_proj": "rowwise"
-        # replicating shared experts (keep them commented out)
-        # "shared_expert_gate_proj": "colwise"
-        # "shared_expert_up_proj": "colwise"
-        # "shared_expert_down_proj": "rowwise"
-        # gating layer should be replicated as well
-        # "gate": "gather"
diff --git a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml
index 12175d03d673..d6462d1ca52d 100644
--- a/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml
+++ b/examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml
@@ -16,22 +16,3 @@ model_kwargs:
 transforms:
   export_to_gm:
     num_moe_experts_for_export: 2
-  detect_sharding:
-    sharding_dims: ['tp','ep', 'bmm']
-    # use only manual config for TP sharding
-    sharding_source: ['manual']
-    manual_config:
-      tp_plan:
-        # GDN layer
-        "in_proj_qkv": "delta"
-        # attention layer
-        "q_proj": "colwise"
-        "k_proj": "colwise"
-        "v_proj": "colwise"
-        "o_proj": "rowwise"
-        # replicating shared experts (keep them commented out)
-        # "shared_expert_gate_proj": "colwise"
-        # "shared_expert_up_proj": "colwise"
-        # "shared_expert_down_proj": "rowwise"
-        # gating layer should be replicated as well
-        # "gate": "gather"
diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
index bc56a877a608..fc780bcef183 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
@@ -189,6 +189,7 @@ def _init_mapping(self):
                 moe_tp_size=self.dist_mapping.get("moe_tp", 1),
                 moe_ep_size=self.dist_mapping.get("moe_ep", self.world_size),
                 moe_cluster_size=self.dist_mapping.get("moe_cluster", 1),
+                enable_attention_dp=self.enable_attention_dp,
             )
         except ValueError as e:
             ad_logger.warning(f"Invalid parallel grid config: {e}")
@@ -239,9 +240,7 @@ def validate_config(self, sources: Union[ShardingSource, List[ShardingSource]] =
             supported_modes = {
                 "colwise",  # row split and no collective
                 "rowwise",  # column split and all-reduce
-                "mla",  # mla layer
                 "mamba",  # mamba SSM layer
-                "delta",  # gated delta net layer
                 "gather",  # simple shard (row + all_gather)
                 # TODO: remaining values are not supported yet.
                 # They require hybrid EP+TP and/or SP support.
@@ -1360,9 +1359,9 @@ def _shard_parameter_node(
 
     # Shard weight using the unified function (also updates the parameter)
     weight_nodes = extract_weight_nodes(node)
-    # Parametrized nodes must have at least one weight (for debugging)
-    assert len(weight_nodes.weights) > 0, (
-        f"Node {node.name} has no weights - weight mapping may be incorrect"
+    # Parametrized nodes must have at least one weight or bias (for debugging)
+    assert len(weight_nodes.weights) > 0 or len(weight_nodes.biases) > 0, (
+        f"Node {node.name} has no weights or biases - weight mapping may be incorrect"
     )
 
     for weight_node in weight_nodes.weights:
@@ -1727,7 +1726,7 @@ def _insert_sharded_mxfp4_mlp_ep(
 def _process_simple_shard(
     nodes_linear: Dict[Node, List[Node]],
     transform_container: ShardingTransformContainer,
-    layer_type: LayerType = LayerType.MLP,
+    layer_type: LayerType = LayerType.UNKNOWN,
 ) -> int:
     # for every linear node:
     # --> row_split (dim 0 of weight) + all_gather (dim -1 of output)
@@ -2665,8 +2664,6 @@ def detect_sharding_from_config(
     num_row_col_shards = 0
     num_attention_shards = 0
     num_ssm_shards = 0
-    num_mla_shards = 0
-    num_delta_shards = 0
     linear_nodes = list(filtered_nodes(gm.graph.nodes, is_any_lin_op))
 
     # use layer_subgraphs to determine the layer_type
@@ -2728,14 +2725,7 @@ def detect_sharding_from_config(
                     if _process_ssm_sharding(layer_subgraph, transform_container) > 0:
                         num_ssm_shards += 1
                         num_row_col_shards += 1
-                elif config == "delta":
-                    if _process_delta_sharding(layer_subgraph, transform_container) > 0:
-                        num_delta_shards += 1
-                        num_row_col_shards += 1
-                elif config == "mla":
-                    if _process_mla_sharding(layer_subgraph, transform_container) > 0:
-                        num_mla_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.")
@@ -2794,8 +2784,7 @@ def detect_sharding_from_config(
     num_shards = num_simple_shards + num_row_col_shards
     ad_logger.info(
         f"Applied {num_shards} TP shards from config. Simple: {num_simple_shards}, "
-        f"row-col: {num_row_col_shards} (attention: {num_attention_shards}, "
-        f"mla: {num_mla_shards}, delta: {num_delta_shards}, ssm: {num_ssm_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)
diff --git a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
index 22ef6c26be2f..b57249b0300f 100644
--- a/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
+++ b/tensorrt_llm/_torch/auto_deploy/utils/node_utils.py
@@ -413,11 +413,18 @@ def find_get_attr_node(weight_node: Node) -> Node:
 
 
 def get_weight_node(node: Node) -> Node:
-    """Get the primary weight node for a compute node."""
+    """Get the primary weight node for a compute node.
+
+    When the node itself is a bias get_attr node (i.e. extract_weight_nodes
+    puts it into .biases rather than .weights), return the bias node so that
+    num_users_of_weight_node gives the correct user count instead of 0.
+    """
     weight_nodes = extract_weight_nodes(node)
-    if len(weight_nodes.weights) == 0:
-        raise ValueError(f"Node {node.name} has no weight")
-    return weight_nodes.weights[0].node
+    if len(weight_nodes.weights) > 0:
+        return weight_nodes.weights[0].node
+    if len(weight_nodes.biases) > 0:
+        return weight_nodes.biases[0].node
+    raise ValueError(f"Node {node.name} has no weight or bias")
 
 
 def num_users_of_weight_node(node: Node) -> int:
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 38a6ec03ed6c..03765ea0386f 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
@@ -18,12 +18,13 @@
 )
 from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer
 from tensorrt_llm._torch.auto_deploy.utils._graph import lint, recompile
+from tensorrt_llm._torch.auto_deploy.utils.mapping_utils import deserialize_mapping
 from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op
 from tensorrt_llm.functional import AllReduceStrategy
 
 
 def _run_ep_shard_job(
-    num_experts: int, rank: int, world_size: int, enable_attention_dp: bool
+    num_experts: int, enable_attention_dp: bool, rank: int, world_size: int
 ) -> None:
     device = "cuda"
     hidden_size = 32
@@ -51,7 +52,7 @@ def _get_expected_num_params(rank: int, world_size: int, num_p_og: int) -> int:
         {
             "detect_sharding": {
                 "stage": "sharding",
-                "sharding_dims": ["ep"],
+                "sharding_dims": ["tp", "ep"],
                 "enable_attention_dp": enable_attention_dp,
             },
             "sharding_transform_executor": {
@@ -68,7 +69,13 @@ def transform_check(gm):
             moe_nodes = [n for n in gm.graph.nodes if is_op(n, torch.ops.auto_deploy.torch_moe)]
             assert len(moe_nodes) == 1, f"Expected 1 MoE node, got {len(moe_nodes)}"
             moe_node = moe_nodes[0]
-            return moe_node.kwargs["enable_alltoall"] == (world_size > 1)
+            print(f"\nMoE node {moe_node.name}, kwargs: {moe_node.kwargs}\n\n")
+            assert "mapping_config" in moe_node.kwargs, (
+                f"Mapping config not found in MoE node {moe_node.name}"
+            )
+            # deserialize the mapping config string to dict
+            mapping_config = deserialize_mapping(moe_node.kwargs["mapping_config"])
+            return mapping_config.enable_attention_dp
     else:
 
         def transform_check(gm):
@@ -76,6 +83,9 @@ def transform_check(gm):
                 is_op(n, torch.ops.auto_deploy.torch_dist_all_reduce) for n in gm.graph.nodes
             ) == (world_size > 1)
 
+    # NOTE: we need to skip output assert when enable_attention_dp is True because
+    # a single rank does not contain the full output (lack of final collective),
+    # so the output will be inherently incomplete w.r.t to the single node.
     run_test_transformed_gm(
         model,
         x,
@@ -83,6 +93,7 @@ def transform_check(gm):
         check_transformed_graph=transform_check,
         _get_expected_num_params=partial(_get_expected_num_params, rank, world_size),
         test_load_hook=False,
+        skip_output_assert=enable_attention_dp,
     )
 
 
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 37ff829e837a..e2d649daf410 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
@@ -26,33 +26,6 @@
 from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_linear_op, is_op, is_weight_node
 from tensorrt_llm.functional import AllReduceStrategy
 
-base_model_tp_plan = {
-    # ATTENTION
-    "q_proj": "colwise",
-    "k_proj": "colwise",
-    "v_proj": "colwise",
-    "o_proj": "rowwise",
-    # gated MLP
-    "gate_proj": "colwise",
-    "up_proj": "colwise",
-    "down_proj": "rowwise",
-    # GPT-style MLP
-    "linear1": "colwise",
-    "linear2": "rowwise",
-    "linear": "gather",
-    # Mamba2 specific projections
-    "in_proj": "mamba",
-    # MLA specific projections
-    "q_a_proj": "mla",
-    # delta net
-    # qkv regex matches both fused (qkvz) and unfused (qkv) opening nodes
-    "in_proj_qkv": "delta",
-}
-
-predefined_config = {
-    "tp_plan": base_model_tp_plan,
-}
-
 
 class GQA_Block(nn.Module):
     def __init__(
@@ -407,6 +380,7 @@ def _run_sharding_execution_job(
     rank: int,
     world_size: int,
 ) -> None:
+    # global predefined_config
     # init model and input
     batch_size = 4
     sequence_len = 8
@@ -423,8 +397,27 @@ def _run_sharding_execution_job(
             hidden_size=num_features,
             num_key_value_heads=num_key_value_heads,
         ).to(device="cuda", dtype=torch.float16)
+        predefined_config = {
+            "tp_plan": {
+                "q_proj": "colwise",
+                "k_proj": "colwise",
+                "v_proj": "colwise",
+                "o_proj": "rowwise",
+            }
+        }
     elif model_cls == FP8MLP:
         model = model_cls(num_features, num_features, bias=bias).to("cuda")
+        predefined_config = {
+            "tp_plan": {
+                # gated MLP
+                "gate_proj": "colwise",
+                "up_proj": "colwise",
+                "down_proj": "rowwise",
+                # GPT-style MLP
+                "linear1": "colwise",
+                "linear2": "rowwise",
+            }
+        }
     elif model_cls == NemotronHMamba2Mixer:
         # Create config for Mamba2 based on Nemotron models
         # Scaled down from typical values: hidden_size=5120, ssm_state_size=128
@@ -450,6 +443,7 @@ def _run_sharding_execution_job(
             num_hidden_layers=1,
         )
         model = model_cls(mamba_config, layer_idx=0).to(device="cuda", dtype=torch.float16)
+        predefined_config = {"tp_plan": {"in_proj": "mamba"}}
     elif model_cls == MLA_Block:
         # Use actual DeepSeek-V3/R1 production values
         # From HuggingFace config (HunYuanPretrainedConfig defaults):
@@ -473,6 +467,7 @@ def _run_sharding_execution_job(
             v_head_dim=v_head_dim,
             bias=bias,
         ).to(device="cuda", dtype=torch.float16)
+        predefined_config = {"tp_plan": {"q_a_proj": "mla"}}
     elif model_cls in (GDN_Block, GDN_Block_Unfused):
         num_k_heads_gdn = 4
         num_v_heads_gdn = 8
@@ -487,10 +482,28 @@ def _run_sharding_execution_job(
             head_v_dim=head_v_dim,
             bias=bias,
         ).to(device="cuda", dtype=torch.float16)
+        predefined_config = {"tp_plan": {"in_proj_qkv": "delta"}}
+    elif model_cls == nn.Linear:
+        model = model_cls(num_features, num_features, bias=bias).to(
+            device="cuda", dtype=torch.float16
+        )
+        # update the tp_plan in predefined_config to force simple sharding of the single linear layer
+        predefined_config = {"tp_plan": {"*": "gather"}}
     else:
         model = model_cls(num_features, num_features, bias=bias).to(
             device="cuda", dtype=torch.float16
         )
+        predefined_config = {
+            "tp_plan": {
+                # gated MLP
+                "gate_proj": "colwise",
+                "up_proj": "colwise",
+                "down_proj": "rowwise",
+                # GPT-style MLP
+                "linear1": "colwise",
+                "linear2": "rowwise",
+            }
+        }
     x = torch.randn(batch_size, sequence_len, num_features, device="cuda", dtype=torch.float16)
 
     # Scale down GDN_Block weights if needed to avoid numerical instability
@@ -531,10 +544,15 @@ def _get_expected_num_params(num_p_og: int) -> int:
         else:
             num_params = num_p_og // world_size + num_update
         if model_cls == MLA_Block:
-            # since q_a_proj is simple-sharded and followed by q_a_layernorm, the layernorm params
-            # are NOT sharded - they have to be replicated. To account for this, we need to add the
-            # number of parameters of the layernorm (weight and bias)to the number of parameters of the model.
-            num_params += 2 * kv_lora_rank * (world_size - 1) // world_size
+            # Both q_a_layernorm and kv_a_layernorm are replicated because their preceding
+            # projections (q_a_proj, kv_a_proj_with_mqa) use simple-shard/all_gather.
+            # Each LayerNorm has weight + bias of size kv_lora_rank → 4 * kv_lora_rank
+            # replicated params total.  The base formula (num_p_og // world_size) treated
+            # these as sharded, so we add back the under-counted fraction.
+            replicated_ln_params = (
+                4 * kv_lora_rank
+            )  # q_a_layernorm + kv_a_layernorm, weight+bias each
+            num_params += replicated_ln_params * (world_size - 1) // world_size
         return num_params
 
     def verify_local_weight_sizes(gm) -> bool:
@@ -563,6 +581,8 @@ def verify_local_weight_sizes(gm) -> bool:
                 "sharding_scope": ["tp"],
                 "sharding_source": [sharding_source],
                 "manual_config": predefined_config,
+                # needed for a solitary Linear node that is not part of any layer subgraph
+                "shard_all_unprocessed": True,
             },
             "sharding_transform_executor": {
                 "stage": "sharding",
@@ -610,6 +630,14 @@ def _run_pattern_detection_job(
             hidden_size=num_features,
             num_key_value_heads=num_key_value_heads,
         ).to(device="cuda", dtype=torch.float16)
+        predefined_config = {
+            "tp_plan": {
+                "q_proj": "colwise",
+                "k_proj": "colwise",
+                "v_proj": "colwise",
+                "o_proj": "rowwise",
+            }
+        }
     elif model_cls == NemotronHMamba2Mixer:
         # Create config for Mamba2
         mamba_config = SimpleNamespace(
@@ -634,6 +662,7 @@ def _run_pattern_detection_job(
             num_hidden_layers=1,
         )
         model = model_cls(mamba_config, layer_idx=0).to(device="cuda", dtype=torch.float16)
+        predefined_config = {"tp_plan": {"in_proj": "mamba"}}
     elif model_cls == MLA_Block:
         # Create simplified MLA based on DeepSeek-V3 architecture
         qk_nope_head_dim = 2
@@ -651,6 +680,7 @@ def _run_pattern_detection_job(
             v_head_dim=v_head_dim,
             bias=bias,
         ).to(device="cuda", dtype=torch.float16)
+        predefined_config = {"tp_plan": {"q_a_proj": "mla"}}
     elif model_cls in (GDN_Block, GDN_Block_Unfused):
         num_k_heads_gdn = 4
         num_v_heads_gdn = 8
@@ -664,10 +694,26 @@ def _run_pattern_detection_job(
             head_v_dim=head_v_dim,
             bias=bias,
         ).to(device="cuda", dtype=torch.float16)
+        predefined_config = {"tp_plan": {"in_proj_qkv": "delta"}}
+    elif model_cls == nn.Linear:
+        model = model_cls(num_features, num_features, bias=bias).to(
+            device="cuda", dtype=torch.float16
+        )
+        # update the tp_plan in predefined_config to force simple sharding of the single linear layer
+        predefined_config = {"tp_plan": {"*": "gather"}}
     else:
         model = model_cls(num_features, num_features, bias=bias).to(
             device="cuda", dtype=torch.float16
         )
+        predefined_config = {
+            "tp_plan": {
+                "gate_proj": "colwise",
+                "up_proj": "colwise",
+                "down_proj": "rowwise",
+                "linear1": "colwise",
+                "linear2": "rowwise",
+            }
+        }
     x = torch.randn(batch_size, sequence_len, num_features, device="cuda", dtype=torch.float16)
 
     # Test pattern detection - create expected transformations for validation
@@ -739,7 +785,7 @@ def _run_pattern_detection_job(
                             config=config,
                             dist_op="all_gather",
                             min_local_shape=1,
-                            layer_type=LayerType.MLP,
+                            layer_type=LayerType.UNKNOWN,
                         )
                     )
         elif model_cls == FP8MLP:

From 69503acc93aaf5745fa8119ac0cb8be315873128 Mon Sep 17 00:00:00 2001
From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Date: Sat, 21 Feb 2026 23:54:54 -0800
Subject: [PATCH 6/7] fix

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
---
 .../auto_deploy/transform/library/sharding.py      | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
index fc780bcef183..3acbdeaf74ee 100644
--- a/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
+++ b/tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py
@@ -2664,6 +2664,8 @@ def detect_sharding_from_config(
     num_row_col_shards = 0
     num_attention_shards = 0
     num_ssm_shards = 0
+    num_mla_shards = 0
+    num_delta_shards = 0
     linear_nodes = list(filtered_nodes(gm.graph.nodes, is_any_lin_op))
 
     # use layer_subgraphs to determine the layer_type
@@ -2725,7 +2727,14 @@ def detect_sharding_from_config(
                     if _process_ssm_sharding(layer_subgraph, transform_container) > 0:
                         num_ssm_shards += 1
                         num_row_col_shards += 1
-
+                elif config == "delta":
+                    if _process_delta_sharding(layer_subgraph, transform_container) > 0:
+                        num_delta_shards += 1
+                        num_row_col_shards += 1
+                elif config == "mla":
+                    if _process_mla_sharding(layer_subgraph, transform_container) > 0:
+                        num_mla_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.")
@@ -2784,7 +2793,8 @@ def detect_sharding_from_config(
     num_shards = num_simple_shards + num_row_col_shards
     ad_logger.info(
         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})"
+        f"row-col: {num_row_col_shards} (attention: {num_attention_shards}, "
+        f"mla: {num_mla_shards}, delta: {num_delta_shards}, ssm: {num_ssm_shards})"
     )
 
     num_matches = len(transform_container.weight_sharding_transforms)

From d3cbd7ffe5cd51bda2dc0a8f9fe9af331bef8a31 Mon Sep 17 00:00:00 2001
From: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
Date: Sun, 22 Feb 2026 23:13:46 -0800
Subject: [PATCH 7/7] refactor

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>
---
 tensorrt_llm/_torch/auto_deploy/distributed/common.py       | 6 ++++--
 .../multigpu/transformations/library/test_ep_sharding.py    | 2 +-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/auto_deploy/distributed/common.py b/tensorrt_llm/_torch/auto_deploy/distributed/common.py
index 4a92c6ee4f38..aba29b21dce2 100644
--- a/tensorrt_llm/_torch/auto_deploy/distributed/common.py
+++ b/tensorrt_llm/_torch/auto_deploy/distributed/common.py
@@ -16,6 +16,8 @@
 
 # TODO: check to what extend we can reuse _torch/distributed.py
 
+_MASTER_ADDR = "127.0.0.1"
+
 
 class _DistGroup:
     """Global instance to set/get the default process group for distributed ops."""
@@ -125,7 +127,7 @@ def _set_distributed_env_vars(local_rank: int, world_size: int, port: int) -> No
     """Set environment variables required by NCCL's env:// init method."""
     os.environ["RANK"] = str(local_rank)
     os.environ["WORLD_SIZE"] = str(world_size)
-    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_ADDR"] = _MASTER_ADDR
     os.environ["MASTER_PORT"] = str(port)
     os.environ["LOCAL_RANK"] = str(local_rank)
 
@@ -135,7 +137,7 @@ def _is_port_available(port: int) -> bool:
     try:
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
             s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-            s.bind(("127.0.0.1", port))
+            s.bind((_MASTER_ADDR, port))
             return True
     except OSError:
         return False
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 03765ea0386f..863cc56625d9 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
@@ -52,7 +52,7 @@ def _get_expected_num_params(rank: int, world_size: int, num_p_og: int) -> int:
         {
             "detect_sharding": {
                 "stage": "sharding",
-                "sharding_dims": ["tp", "ep"],
+                "sharding_dims": ["ep"],
                 "enable_attention_dp": enable_attention_dp,
             },
             "sharding_transform_executor": {