Fixed QKVO col-row sharding

Signed-off-by: greg-kwasniewski1 <213329731+greg-kwasniewski1@users.noreply.github.com>

Author: greg-kwasniewski1

tensorrt_llm/_torch/auto_deploy/models/patches/nemotron_h.py

@@ -148,9 +148,13 @@ def get_model_from_config_patched(config, **kwargs):
 _nemotron_h_base_model_tp_plan = {
     "in_proj": "mamba",
     "out_proj": "rowwise",
+    "q_proj": "colwise",
+    "k_proj": "colwise",
+    "v_proj": "colwise",
+    "o_proj": "rowwise",
     "up_proj": "colwise",
     "down_proj": "rowwise",
-    "*": "gather",
+    # "*": "gather",
 }
 
 

tensorrt_llm/_torch/auto_deploy/utils/sharding_utils.py

@@ -83,8 +83,8 @@ def _validate_sharded_shapes(
     next_lin_node, _ = bfs(node, is_any_lin_op, include_root=False)
     nodes_to_validate = subgraph(
         [node],
-        [next_lin_node],
         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:
@@ -96,7 +96,7 @@ def _validate_sharded_shapes(
             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  # view_shape[2] // world_size
             args[1] = tuple(view_shape)
             view_node.args = tuple(args)
             view_node.meta["sharded"] = True