feat: Add pp support for hybrid attn/mamba model

examples/pytorch/out_of_tree_example/modeling_opt.py

@@ -89,6 +89,7 @@ def __init__(
 
     def forward(
         self,
+        position_ids: torch.LongTensor,
         hidden_states: torch.Tensor,
         attn_metadata: AttentionMetadata,
         **kwargs,
@@ -211,8 +212,11 @@ def forward(
 
         # residual = None
         for decoder_layer in self.layers:
-            hidden_states = decoder_layer(hidden_states=hidden_states,
-                                          attn_metadata=attn_metadata)
+            hidden_states = decoder_layer(
+                position_ids=position_ids,
+                hidden_states=hidden_states,
+                attn_metadata=attn_metadata,
+            )
 
         if self.final_layer_norm is not None:
             hidden_states = self.final_layer_norm(hidden_states)

tensorrt_llm/_torch/models/modeling_mamba_hybrid.py

@@ -10,6 +10,7 @@
 from ..attention_backend import AttentionMetadata
 from ..model_config import ModelConfig
 from ..modules.attention import Attention
+from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
 from ..modules.mamba import Mamba2, MambaCacheManager, MambaCacheParams
 from ..modules.mlp import MLP
@@ -93,7 +94,7 @@ def forward(
                                attn_metadata=attn_metadata)
 
 
-class MambaHybridLayer(nn.Module):
+class MambaHybridLayer(DecoderLayer):
 
     def __init__(
         self,
@@ -138,8 +139,10 @@ def __init__(
 
     def forward(
         self,
+        position_ids: torch.LongTensor,
         hidden_states: torch.Tensor,
         attn_metadata: AttentionMetadata,
+        *,
         mamba_cache_params: MambaCacheParams,
         residual: Optional[torch.Tensor] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -240,8 +243,13 @@ def forward(
         hidden_states = inputs_embeds
         residual = None
         for layer in self.layers:
-            hidden_states, residual = layer(hidden_states, attn_metadata,
-                                            mamba_cache_params, residual)
+            hidden_states, residual = layer(
+                position_ids,
+                hidden_states,
+                attn_metadata,
+                mamba_cache_params=mamba_cache_params,
+                residual=residual,
+            )
 
         hidden_states, _ = self.norm_f(hidden_states, residual)
 

tensorrt_llm/_torch/models/modeling_nemotron_h.py

@@ -13,6 +13,7 @@
 from ..attention_backend import AttentionMetadata
 from ..model_config import ModelConfig
 from ..modules.attention import Attention
+from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
 from ..modules.mamba.mixer import MambaMixer
 from ..modules.mlp import MLP
@@ -93,7 +94,7 @@ def forward(
                                attn_metadata=attn_metadata)
 
 
-class NemotronHLayer(nn.Module):
+class NemotronHLayer(DecoderLayer):
 
     def __init__(
         self,
@@ -140,6 +141,7 @@ def __init__(
 
     def forward(
         self,
+        position_ids: torch.LongTensor,
         hidden_states: torch.Tensor,
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
@@ -167,23 +169,8 @@ def __init__(self, model_config: ModelConfig[NemotronHConfig]):
         )
 
         # create layers
-        mlp_idx = 0
-        mamba_idx = 0
-        transformer_idx = 0
         layers = []
-        for layer_type in config.hybrid_override_pattern:
-            # calculate layer index based on type
-            if layer_type == "M":
-                layer_idx = mamba_idx
-                mamba_idx += 1
-            elif layer_type == "-":
-                layer_idx = mlp_idx
-                mlp_idx += 1
-            elif layer_type == "*":
-                layer_idx = transformer_idx
-                transformer_idx += 1
-            else:
-                ValueError(f"{layer_type} is not supported")
+        for layer_idx, layer_type in enumerate(config.hybrid_override_pattern):
             layers.append(NemotronHLayer(model_config, layer_idx, layer_type))
         self.layers = nn.ModuleList(layers)
 
@@ -215,7 +202,7 @@ def forward(
         hidden_states = inputs_embeds
 
         for layer in self.layers:
-            hidden_states = layer(hidden_states, attn_metadata)
+            hidden_states = layer(position_ids, hidden_states, attn_metadata)
 
         hidden_states = self.norm_f(hidden_states)
 

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -257,7 +257,13 @@ def create_kv_cache_manager(model_engine: PyTorchModelEngine, mapping: Mapping,
         elif is_nemotron_hybrid(config):
             config = model_engine.model.model_config.pretrained_config
             num_layers = config.hybrid_override_pattern.count("*")
+            layer_mask = [
+                char == "*" for char in config.hybrid_override_pattern
+            ]
             mamba_num_layers = config.hybrid_override_pattern.count("M")
+            mamba_layer_mask = [
+                char == "M" for char in config.hybrid_override_pattern
+            ]
             kv_cache_manager = MambaHybridCacheManager(
                 # mamba cache parameters
                 config.hidden_size,
@@ -267,11 +273,13 @@ def create_kv_cache_manager(model_engine: PyTorchModelEngine, mapping: Mapping,
                 config.n_groups,
                 config.mamba_head_dim,
                 mamba_num_layers,
+                mamba_layer_mask,
                 config.torch_dtype,
                 # kv cache parameters
                 executor_config.kv_cache_config,
                 tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF,
                 num_layers=num_layers,
+                layer_mask=layer_mask,
                 num_kv_heads=num_key_value_heads,
                 head_dim=head_dim,
                 tokens_per_block=executor_config.tokens_per_block,

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -67,19 +67,30 @@ def get_pp_layers(
     num_layers: int,
     mapping: Mapping,
     spec_config: Optional["SpecConfig"] = None,
+    layer_mask: Optional[List[bool]] = None,
 ) -> Tuple[List[int], int]:
     from ..speculative.utils import get_num_spec_layers
 
-    pp_layers = mapping.pp_layers(num_layers)
+    total_num_layers = num_layers
+    if layer_mask is not None:
+        assert sum(layer_mask) == num_layers, (
+            f"The number of enabled layers in layer_mask ({sum(layer_mask)}) "
+            f"must match the number of layers ({num_layers}) "
+            f"in KV cache manager, but get layer_mask: {layer_mask}")
+        total_num_layers = len(layer_mask)
+    pp_layers = mapping.pp_layers(total_num_layers)
+    if layer_mask is not None:
+        pp_layers = [i for i in pp_layers if layer_mask[i]]
     if spec_config is not None:
         num_spec_layers = get_num_spec_layers(spec_config)
-        num_layers += num_spec_layers
+        total_num_layers += num_spec_layers
         if mapping.is_last_pp_rank():
-            pp_layers.extend(range(num_layers - num_spec_layers, num_layers))
+            pp_layers.extend(
+                range(total_num_layers - num_spec_layers, total_num_layers))
     if len(pp_layers) == 0:
         # Don't support empty KV cache for now, provide at least 1 layer
         pp_layers.append(0)
-    return pp_layers, num_layers
+    return pp_layers, total_num_layers
 
 
 class KVCacheManager(BaseResourceManager):
@@ -100,12 +111,17 @@ def __init__(
         mapping: Mapping,
         dtype: DataType = DataType.HALF,
         spec_config: Optional["SpecConfig"] = None,
+        layer_mask: Optional[List[bool]] = None,
     ) -> None:
         self.mapping = mapping
         self.dtype = dtype
         self.kv_cache_type = kv_cache_type
-        self.pp_layers, self.num_layers = get_pp_layers(num_layers, mapping,
-                                                        spec_config)
+        self.pp_layers, self.num_layers = get_pp_layers(
+            num_layers,
+            mapping,
+            spec_config=spec_config,
+            layer_mask=layer_mask,
+        )
         self.num_local_layers = len(self.pp_layers)
         self.layer_offsets = {
             idx: offset
@@ -127,9 +143,8 @@ def __init__(
 
             self.num_kv_heads_per_layer = []
             if self.num_local_layers > 0:
-                for kv_head in num_kv_heads[self.
-                                            pp_layers[0]:self.pp_layers[-1] +
-                                            1]:
+                for i in self.pp_layers:
+                    kv_head = num_kv_heads[i]
                     if kv_head is not None:
                         self.num_kv_heads_per_layer.append(
                             (kv_head + tp_size - 1) // tp_size)
@@ -485,10 +500,20 @@ def rewind_kv_cache(self, request: LlmRequest, rewind_len: int):
 
 class MambaCacheManager(BaseResourceManager):
 
-    def __init__(self, d_model: int, d_state: int, d_conv: int, expand: int,
-                 n_groups: int, head_dim: int, num_layers: int,
-                 max_batch_size: int, mapping: Mapping,
-                 conv1d_state_dtype: torch.dtype) -> None:
+    def __init__(
+        self,
+        d_model: int,
+        d_state: int,
+        d_conv: int,
+        expand: int,
+        n_groups: int,
+        head_dim: int,
+        num_layers: int,
+        max_batch_size: int,
+        mapping: Mapping,
+        conv1d_state_dtype: torch.dtype,
+        layer_mask: Optional[List[bool]] = None,
+    ) -> None:
 
         # get tp size
         tp_size = mapping.tp_size
@@ -509,7 +534,11 @@ def __init__(self, d_model: int, d_state: int, d_conv: int, expand: int,
         # conv and ssm states device
         device = torch.device("cuda")
 
-        pp_layers, num_layers = get_pp_layers(num_layers, mapping)
+        pp_layers, num_layers = get_pp_layers(
+            num_layers,
+            mapping,
+            layer_mask=layer_mask,
+        )
         num_local_layers = len(pp_layers)
         self.mamba_layer_offsets = {
             idx: offset
@@ -607,12 +636,14 @@ def __init__(
         mamba_n_groups: int,
         mamba_head_dim: int,
         mamba_num_layers: int,
+        mamba_layer_mask: List[bool],
         mamba_conv1d_state_dtype: torch.dtype,
         # kv cache parameters
         kv_cache_config: KvCacheConfigCpp,
         kv_cache_type: CacheTypeCpp,
         *,
         num_layers: int,
+        layer_mask: List[bool],
         num_kv_heads: Union[int, List[Optional[int]]],
         head_dim: int,
         tokens_per_block: int,
@@ -629,25 +660,37 @@ def __init__(
         assert not kv_cache_config.enable_block_reuse, "mamba hybrid cache requires block reuse to be disabled in KV cache config"
 
         # initialize mamba cache manager
-        MambaCacheManager.__init__(self, mamba_d_model, mamba_d_state,
-                                   mamba_d_conv, mamba_expand, mamba_n_groups,
-                                   mamba_head_dim, mamba_num_layers,
-                                   max_batch_size, mapping,
-                                   mamba_conv1d_state_dtype)
+        MambaCacheManager.__init__(
+            self,
+            mamba_d_model,
+            mamba_d_state,
+            mamba_d_conv,
+            mamba_expand,
+            mamba_n_groups,
+            mamba_head_dim,
+            mamba_num_layers,
+            max_batch_size,
+            mapping,
+            mamba_conv1d_state_dtype,
+            mamba_layer_mask,
+        )
 
         # initialize kv cache manager
-        KVCacheManager.__init__(self,
-                                kv_cache_config,
-                                kv_cache_type,
-                                num_layers=num_layers,
-                                num_kv_heads=num_kv_heads,
-                                head_dim=head_dim,
-                                tokens_per_block=tokens_per_block,
-                                max_seq_len=max_seq_len,
-                                max_batch_size=max_batch_size,
-                                mapping=mapping,
-                                dtype=dtype,
-                                spec_config=spec_config)
+        KVCacheManager.__init__(
+            self,
+            kv_cache_config,
+            kv_cache_type,
+            num_layers=num_layers,
+            num_kv_heads=num_kv_heads,
+            head_dim=head_dim,
+            tokens_per_block=tokens_per_block,
+            max_seq_len=max_seq_len,
+            max_batch_size=max_batch_size,
+            mapping=mapping,
+            dtype=dtype,
+            spec_config=spec_config,
+            layer_mask=layer_mask,
+        )
 
     def prepare_resources(self, scheduled_batch: ScheduledRequests):
         self.prepare_mamba_resources(scheduled_batch)

tests/unittest/_torch/modeling/test_modeling_nemotron_h.py

@@ -92,7 +92,13 @@ def test_nemotron_h_sanity(self):
         tokens_per_block = 128
         head_dim = nemotron_h.config.hidden_size // nemotron_h.config.num_attention_heads
         num_layers = nemotron_h.config.hybrid_override_pattern.count("*")
+        layer_mask = [
+            char == "*" for char in nemotron_h.config.hybrid_override_pattern
+        ]
         mamba_num_layers = nemotron_h.config.hybrid_override_pattern.count("M")
+        mamba_layer_mask = [
+            char == "M" for char in nemotron_h.config.hybrid_override_pattern
+        ]
         num_kv_heads = nemotron_h.config.num_key_value_heads
         max_seq_len = num_blocks * tokens_per_block
         batch_size = len(context_sequence_lengths) + 2
@@ -115,11 +121,13 @@ def test_nemotron_h_sanity(self):
             nemotron_h.config.n_groups,
             nemotron_h.config.mamba_head_dim,
             mamba_num_layers,
+            mamba_layer_mask,
             nemotron_h.config.torch_dtype,
             # kv cache parameters
             kv_cache_config,
             tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF,
             num_layers=num_layers,
+            layer_mask=layer_mask,
             num_kv_heads=num_kv_heads,
             head_dim=head_dim,
             tokens_per_block=tokens_per_block,