[Feature][AutoDeploy]: Heuristics to shard Mamba layers

[greg-kwasniewski1]

🚀 The feature, motivation and pitch

Currently, we shard Mamba layers only through the predefined config and "one large hardcoded transformation". TP sharding heuristics should properly detect and shard Mamba layers.

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

[coderabbitai]

A summary of the changes CodeRabbit can apply:

• Implement automatic Mamba layer sharding by modifying tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py: add Mamba operation recognition sets (shardable_mamba_nodes, shardable_nodes_with_mamba), introduce mamba_nodes and mamba_related_nodes tracking, update node classification to detect Mamba ops, add post-processing to clear related nodes, include them in unaccounted-node checks, and extract head-dimension metadata to enforce correct column/row sharding and all-reduce behavior.

• Update tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py to add Mamba/SSM-aware column/row sharding: introduce shardable_mamba and related node sets, track mamba and mamba-related nodes alongside attention nodes, extend heuristics to clear related nodes and fall back to simple sharding when ambiguous, and extract head_dim from Mamba node outputs to apply column-parallel in_proj and row-parallel out_proj sharding (with all_reduce) for Mamba layers.

• ✅ Create PR with these edits
• 📋 Get copyable edits