[None][feat] AutoDeploy: Qwen3.5: Apply whielist based sharding and apply lm_head sharding

examples/auto_deploy/model_registry/configs/qwen3.5_moe_400b.yaml

@@ -29,10 +29,16 @@ transforms:
     enabled: true
   fuse_nvfp4_moe:
     backend: trtllm_gen
+  detect_sharding:
+    enabled: false
+  sharding_transform_executor:
+    enabled: false
   apply_sharding_hints:
+    enabled: true
     allreduce_strategy: SYMM_MEM
     # Shared expert is excluded from sharding for performance purpose
     shard_layers: ["moe", "delta", "mha"]
+    simple_shard_filter: "lm_head"
   multi_stream_moe:
     stage: compile
     enabled: true

tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py

@@ -635,26 +635,25 @@ def __init__(self, config: Qwen3_5MoeTextConfig, intermediate_size: int):
         self.act_fn = ACT2FN[config.hidden_act]
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Intentionally left untagged: with no ``layer_type`` it defaults to "unknown",
+        # which the ``shard_layers`` inclusion whitelist excludes -> kept replicated.
         gate = torch.ops.auto_deploy.torch_linear_simple(
             x,
             self.gate_proj.weight,
             self.gate_proj.bias,
             tp_mode="colwise",
-            layer_type="shared_expert",
         )
         up = torch.ops.auto_deploy.torch_linear_simple(
             x,
             self.up_proj.weight,
             self.up_proj.bias,
             tp_mode="colwise",
-            layer_type="shared_expert",
         )
         return torch.ops.auto_deploy.torch_linear_simple(
             self.act_fn(gate) * up,
             self.down_proj.weight,
             self.down_proj.bias,
             tp_mode="rowwise",
-            layer_type="shared_expert",
         )
 
 

tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py

@@ -22,7 +22,6 @@
 The SwiGLU pattern is: silu(x @ gate.T) * (x @ up.T) @ down.T
 """
 
-from contextlib import contextmanager
 from typing import Tuple, Type
 
 import torch
@@ -48,7 +47,7 @@
     eliminate_dead_code,
     get_attr_by_name,
 )
-from ...utils.node_utils import extract_op_args, is_op, set_op_args
+from ...utils.node_utils import is_op
 from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern
 from ...utils.quantization_utils import ensure_tma_col_major
 from ..interface import (
@@ -60,42 +59,6 @@
 )
 
 
-def _weight_key(node: Node):
-    """Stable key for a weight arg: the get_attr target FQN (survives node re-creation
-    during pattern replacement), falling back to node identity. ``args[1]`` is the
-    weight for both linear and SwiGLU ops."""
-    w = node.args[1] if len(node.args) > 1 else None
-    if not isinstance(w, Node):
-        return None
-    return w.target if w.op == "get_attr" else w
-
-
-@contextmanager
-def preserve_layer_types(gm: GraphModule, linear_op, fused_op):
-    """Carry the ``layer_type`` hint across a fusion that consumes ``linear_op`` nodes
-    and emits ``fused_op`` nodes (which would otherwise drop the hint).
-
-    Snapshots each source weight's ``layer_type`` before the rewrite, then re-applies
-    it to the fused node keyed by weight, so hint-driven sharding (``shard_layers``)
-    can still classify the fused node. Wrap the matcher's ``patterns.apply`` call.
-    """
-    wmap = {}
-    for n in gm.graph.nodes:
-        if is_op(n, linear_op):
-            [lt] = extract_op_args(n, "layer_type")
-            key = _weight_key(n)
-            if lt is not None and key is not None:
-                wmap[key] = lt
-    yield
-    if not wmap:
-        return
-    for n in gm.graph.nodes:
-        if is_op(n, fused_op):
-            key = _weight_key(n)
-            if key is not None and key in wmap:
-                set_op_args(n, layer_type=wmap[key])
-
-
 def _maybe_to_deepgemm_layout(
     weight: torch.Tensor, scale: torch.Tensor
 ) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -279,12 +242,7 @@ def _apply(
             dummy_args=dummy_args_with_bias,
         )
 
-        with preserve_layer_types(
-            gm,
-            torch.ops.auto_deploy.torch_linear_simple.default,
-            torch.ops.auto_deploy.torch_swiglu_mlp.default,
-        ):
-            num_matches = patterns.apply(gm.graph)
+        num_matches = patterns.apply(gm.graph)
 
         if num_matches > 0:
             gm.recompile()
@@ -596,12 +554,7 @@ def _apply(
             dummy_args=dummy_args,
         )
 
-        with preserve_layer_types(
-            gm,
-            torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear.default,
-            torch.ops.auto_deploy.torch_nvfp4_swiglu_mlp.default,
-        ):
-            num_matches = patterns.apply(gm.graph)
+        num_matches = patterns.apply(gm.graph)
 
         if num_matches > 0:
             gm.recompile()
@@ -882,12 +835,7 @@ def _apply(
             dummy_args=dummy_args,
         )
 
-        with preserve_layer_types(
-            gm,
-            torch.ops.auto_deploy.torch_fake_quant_finegrained_fp8_linear.default,
-            torch.ops.auto_deploy.torch_finegrained_fp8_swiglu_mlp.default,
-        ):
-            num_matches = patterns.apply(gm.graph)
+        num_matches = patterns.apply(gm.graph)
 
         if num_matches > 0:
             gm.recompile()

tensorrt_llm/_torch/auto_deploy/transform/library/sharding_ir.py

@@ -148,6 +148,37 @@ def _fp4_weight_scale_pipeline_cache_spec(
     }
 
 
+def _rowwise_bias_load_hook(state_dict, prefix, *args, rank, param_key):
+    """Always-apply load hook for the rank0-only row-parallel bias: zero on rank != 0.
+
+    Unlike the shape-gated ``_load_hook`` (which only transforms when the loaded
+    shape differs from the sharded shape), the row-parallel bias keeps its full
+    shape -- only its *value* changes. Idempotent (zeroing zeros / keeping rank 0).
+    Takes plain ``rank``/``param_key`` (not a closure) so the pipeline cache can
+    serialize and rebuild it via the importable-hook path.
+    """
+    key = prefix + param_key
+    if key in state_dict and rank != 0:
+        state_dict[key] = torch.zeros_like(state_dict[key])
+
+
+def _replicate_rowwise_bias(bn: WeightNode, rank: int) -> None:
+    """Keep a row-parallel linear's bias on rank 0 only; zero it on the others.
+
+    The row-parallel output is summed across ranks by the trailing ``all_reduce``,
+    so a full bias present on every rank would be added ``world_size`` times. Zeroing
+    it on rank != 0 makes the all_reduce contribute the bias exactly once. The
+    parameter shape is unchanged (unlike the column-parallel bias, which is split),
+    so a dedicated always-apply load hook is used.
+    """
+    new_bias = bn.tensor if rank == 0 else torch.zeros_like(bn.tensor)
+    pname = bn.node_key.rsplit(".", 1)[-1]
+    bn.submod._register_load_state_dict_pre_hook(
+        partial(_rowwise_bias_load_hook, rank=rank, param_key=pname)
+    )
+    setattr(bn.submod, pname, torch.nn.Parameter(new_bias.detach().clone(), requires_grad=False))
+
+
 _SHARDING_HINT_NAMES = frozenset(
     {
         "tp_mode",
@@ -353,6 +384,10 @@ def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int
                     world_size=dc.tp_size,
                     fused_weight_dims=fused,
                 )
+            else:
+                # Row-parallel: the trailing all_reduce would sum a full bias
+                # world_size times. Keep it on rank 0 only so it is added once.
+                _replicate_rowwise_bias(bn, dc.tp_rank)
 
         ad_logger.debug(f"  sharded linear tp_mode={tp_mode}")
         return 1
@@ -595,6 +630,48 @@ def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int
         return 1 if count > 0 else 0
 
 
+@ShardableNode.register(torch.ops.auto_deploy.torch_attention)
+class AttentionSinksShardableNode(ShardableNode):
+    """``torch_attention`` with per-head ``sinks``: shard sinks along the head dim.
+
+    Attention-sink models (e.g. GPT-OSS) add a learnable per-head sink scalar
+    (shape ``[num_heads]``). When the heads are TP-sharded (q/k/v colwise), the
+    ``sinks`` arg must follow the same head split, else the op sees full sinks
+    against the sharded head count. Standard attention (no ``sinks``) is a no-op.
+
+    Gating is handled by the apply loop: attention-DP skips all non-MoE nodes
+    (so attention stays replicated and sinks stays full), and ``shard_layers``
+    gates via the node's ``layer_type`` hint (default ``"mha"``).
+    """
+
+    def apply(self, gm: GraphModule, dc: DistConfig, max_num_tokens: int = 0) -> int:
+        if dc.tp_size <= 1:
+            return 0
+        count = 0
+        for wn in extract_weight_nodes(self.node).weights:
+            # Only the per-head ``sinks`` (1-D) follows the head split; never a 2-D weight.
+            if wn.tensor.dim() != 1:
+                continue
+            shard_weight_tensor(
+                gm=gm,
+                weight_tensor=wn.tensor,
+                param_key=wn.node_key,
+                dim=0,
+                rank=dc.tp_rank,
+                world_size=dc.tp_size,
+            )
+            count += 1
+        if count:
+            ad_logger.debug("  sharded attention sinks along head dim")
+        return 1 if count > 0 else 0
+
+    @classmethod
+    def _strip_node_hints(cls, node: Node) -> bool:
+        # Leave ``torch_attention`` untouched at strip time: its ``layer_type`` is a
+        # benign op default that downstream backend selection reads as-is.
+        return False
+
+
 @ShardableNode.register(*_auto_deploy_ops("torch_rmsnorm_gated", "triton_rmsnorm_gated"))
 class NormShardableNode(ShardableNode):
     """Gated RMSNorm op: shard weight parameter."""
@@ -979,6 +1056,12 @@ class IRShardingConfig(TransformConfig):
         default=None,
         description="When set, only shard nodes whose layer_type hint is in this list.",
     )
+    simple_shard_filter: Optional[str] = Field(
+        default=None,
+        description="Comma-separated weight-name keywords (e.g. 'lm_head'). Matching linears are "
+        "gather-sharded (column split + all_gather) regardless of shard_layers -- used for the "
+        "lm_head vocab projection, which the hint-driven sharder would otherwise replicate.",
+    )
     enable_attention_dp: bool = Field(default=False)
     dist_mapping: dict[str, int] = Field(default_factory=dict)
     dist_config: DistConfig = Field(default_factory=DistConfig)
@@ -1047,46 +1130,54 @@ def _apply_simple_shard(gm: GraphModule, dc: DistConfig) -> int:
     for node in list(gm.graph.nodes):
         if not is_any_lin_op(node):
             continue
-        weight_nodes = extract_weight_nodes(node)
-        if not weight_nodes.weights:
-            continue
-        for wn in weight_nodes.weights:
-            shard_weight_tensor(
-                gm=gm,
-                weight_tensor=wn.tensor,
-                param_key=wn.node_key,
-                dim=SplitDimension.COLUMN,
-                rank=dc.tp_rank,
-                world_size=dc.tp_size,
-            )
-        for bn in weight_nodes.biases:
-            shard_weight_tensor(
-                gm=gm,
-                weight_tensor=bn.tensor,
-                param_key=bn.node_key,
-                dim=SplitDimension.COLUMN,
-                rank=dc.tp_rank,
-                world_size=dc.tp_size,
-            )
-        enable_sharding = ShardableNode.from_node(node)
-        if isinstance(enable_sharding, LinearShardableNode):
-            enable_sharding._shard_scales(
-                gm, dc, weight_nodes, dim=SplitDimension.COLUMN, min_shape=1, fused=None
-            )
-        # torch_dist_all_gather is the demollm backend op; signature is
-        # (tensor, dim=0, sizes=None) — plain torch.distributed all_gather,
-        # no strategy or symm_mem support (use the trtllm backend for those).
-        with gm.graph.inserting_after(node):
-            gather_node = gm.graph.call_function(
-                torch.ops.auto_deploy.torch_dist_all_gather.default,
-                args=(node, -1),
-            )
-            node.replace_all_uses_with(gather_node)
-            gather_node.replace_input_with(gather_node, node)
-        num_updates += 1
+        num_updates += _simple_shard_node(gm, node, dc)
     return num_updates
 
 
+def _simple_shard_node(gm: GraphModule, node: Node, dc: DistConfig) -> int:
+    """Column-split one linear's weight/bias/scale, then all_gather (the "gather" mode).
+
+    Used as the simple-shard fallback for every linear (``simple_shard_only``) and,
+    via ``simple_shard_filter``, to gather-shard specific linears like ``lm_head``
+    (huge vocab projection) that the hint-driven sharder would otherwise replicate.
+    """
+    weight_nodes = extract_weight_nodes(node)
+    if not weight_nodes.weights:
+        return 0
+    for wn in weight_nodes.weights:
+        shard_weight_tensor(
+            gm=gm,
+            weight_tensor=wn.tensor,
+            param_key=wn.node_key,
+            dim=SplitDimension.COLUMN,
+            rank=dc.tp_rank,
+            world_size=dc.tp_size,
+        )
+    for bn in weight_nodes.biases:
+        shard_weight_tensor(
+            gm=gm,
+            weight_tensor=bn.tensor,
+            param_key=bn.node_key,
+            dim=SplitDimension.COLUMN,
+            rank=dc.tp_rank,
+            world_size=dc.tp_size,
+        )
+    sn = ShardableNode.from_node(node)
+    if isinstance(sn, LinearShardableNode):
+        sn._shard_scales(gm, dc, weight_nodes, dim=SplitDimension.COLUMN, min_shape=1, fused=None)
+    # torch_dist_all_gather is the demollm backend op; signature is
+    # (tensor, dim=0, sizes=None) — plain torch.distributed all_gather,
+    # no strategy or symm_mem support (use the trtllm backend for those).
+    with gm.graph.inserting_after(node):
+        gather_node = gm.graph.call_function(
+            torch.ops.auto_deploy.torch_dist_all_gather.default,
+            args=(node, -1),
+        )
+        node.replace_all_uses_with(gather_node)
+        gather_node.replace_input_with(gather_node, node)
+    return 1
+
+
 # =============================================================================
 # Transform classes
 # =============================================================================
@@ -1190,6 +1281,8 @@ def _apply(
             _log_sharding_result(dc, num_updates)
         else:
             shard_layers = self.config.shard_layers
+            ssf = self.config.simple_shard_filter
+            simple_shard_filter = [k.strip() for k in ssf.split(",") if k.strip()] if ssf else None
             num_skipped = 0
             all_dead_nodes = []
 
@@ -1202,6 +1295,14 @@ def _apply(
                         shardable_node, (MoEShardableNode, StackedMoEShardableNode)
                     ):
                         continue
+                    # Gather-shard simple_shard_filter matches (e.g. lm_head) regardless of
+                    # shard_layers: column split + all_gather of the huge vocab projection.
+                    if simple_shard_filter and is_any_lin_op(node):
+                        wnodes = extract_weight_nodes(node)
+                        key = wnodes.weights[0].node_key if wnodes.weights else ""
+                        if any(kw in key for kw in simple_shard_filter):
+                            num_updates += _simple_shard_node(gm, node, dc)
+                            continue
                     if shard_layers is not None:
                         [lt] = extract_op_args(node, "layer_type")
                         if lt is not None and lt not in shard_layers: