Update tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py

Co-authored-by: Patrice Castonguay <55748270+pcastonguay@users.noreply.github.com>
Signed-off-by: Bo Deng <deemod@nvidia.com>

Author: bo-nv

tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py

@@ -314,7 +314,7 @@ def _mamba_tp(ri: RankInfo) -> Tuple[int, int]:
     @staticmethod
     def _build_layer_ptrs(
         pool,
-        layer_offsets: dict,
+        layer_offsets: Dict[int,int],
         overlapping_layers: List[int],
         slot: int,
     ) -> List[int]: