pack more into mamba2 metadata

Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com>

Author: cyang49

vllm/model_executor/layers/mamba/mamba2_metadata.py

@@ -8,12 +8,13 @@
 @dataclass
 class Mamba2Metadata:
     chunk_size: int
+    seq_idx: torch.Tensor
     chunk_indices: torch.Tensor
     chunk_offsets: torch.Tensor
 
 
-def prepare_mamba2_metadata(seq_idx: torch.Tensor,
-                            chunk_size: int) -> Mamba2Metadata:
+def _seq_idx_to_chunk_indices_offsets(seq_idx, chunk_size: int):
+
     # convert seq_idx to chunk indices and offsets
     # - derive the cu_seqlens
     _, cu_seqlens = torch.where(seq_idx.diff())
@@ -43,6 +44,41 @@ def prepare_mamba2_metadata(seq_idx: torch.Tensor,
         chunk_indices[_s:_e] -= p
         chunk_offsets[_s] = s % chunk_size
 
+    return chunk_indices, chunk_offsets
+
+
+def prepare_mamba2_metadata(
+    chunk_size: int,
+    has_prefills: bool,
+    input_ids: torch.Tensor,
+    query_start_loc: torch.Tensor,
+) -> Mamba2Metadata:
+
+    seq_idx = None
+    chunk_indices, chunk_offsets = None, None
+    if has_prefills:
+        seq_idx = torch.zeros_like(input_ids, dtype=torch.int32)
+        for i, (srt,
+                end) in enumerate(zip(
+                    query_start_loc,
+                    query_start_loc[1:],
+                )):
+            seq_idx[srt:end] = i
+        seq_idx.unsqueeze_(0)
+
+        # compute metadata for chunked prefill.
+        # actually this is only needed if there are
+        # initial states, but this is determinable
+        # only from attention metadata yet
+        # unavailable from the top-level model forward.
+        # Rather than complicating things to extract said
+        # metadata, we simply just compute redundently and
+        # will be silently ignored inside the mamba kernels.
+        # if not needed.
+        chunk_indices, chunk_offsets = _seq_idx_to_chunk_indices_offsets(
+            seq_idx, chunk_size)
+
     return Mamba2Metadata(chunk_size=chunk_size,
+                          seq_idx=seq_idx,
                           chunk_indices=chunk_indices,
                           chunk_offsets=chunk_offsets)

vllm/model_executor/layers/mamba/mamba_mixer2.py

@@ -222,7 +222,6 @@ def __init__(self,
                  head_dim: int = 64,
                  rms_norm_eps: float = 1e-5,
                  activation="silu",
-                 chunk_size: int = 256,
                  quant_config: Optional[QuantizationConfig] = None):
         super().__init__()
 
@@ -258,7 +257,6 @@ def __init__(self,
         self.ssm_state_size = ssm_state_size
         self.activation = activation
 
-        self.chunk_size = chunk_size
         self.intermediate_size = intermediate_size
         self.head_dim = head_dim
         self.num_heads = num_heads
@@ -389,8 +387,7 @@ def forward_cuda(
         self,
         hidden_states: torch.Tensor,
         mamba_cache_params: MambaCacheParams,
-        sequence_idx: Optional[torch.Tensor] = None,
-        mamba2_metadata: Optional[Mamba2Metadata] = None,
+        mamba2_metadata: Mamba2Metadata,
     ):
         # For the mamba2 triton kernels to operate in continuous batching,
         # the sequence_idx is needed to be passed in. Also, for the kernels
@@ -400,11 +397,6 @@ def forward_cuda(
         # layers.
         attn_metadata: AttentionMetadata = get_forward_context().attn_metadata
 
-        chunk_indices, chunk_offsets = None, None
-        if mamba2_metadata is not None:
-            chunk_indices = mamba2_metadata.chunk_indices
-            chunk_offsets = mamba2_metadata.chunk_offsets
-
         seq_len, _ = hidden_states.shape
         groups_time_state_size = self.n_groups * self.ssm_state_size
 
@@ -496,13 +488,13 @@ def forward_cuda(
                 self.A,
                 B.view(1, seq_len, self.n_groups // self.tp_size, -1),
                 C.view(1, seq_len, self.n_groups // self.tp_size, -1),
-                chunk_size=self.chunk_size,
+                chunk_size=mamba2_metadata.chunk_size,
                 D=self.D,
                 z=None,
                 dt_bias=self.dt_bias,
-                seq_idx=sequence_idx,
-                chunk_indices=chunk_indices,
-                chunk_offsets=chunk_offsets,
+                seq_idx=mamba2_metadata.seq_idx,
+                chunk_indices=mamba2_metadata.chunk_indices,
+                chunk_offsets=mamba2_metadata.chunk_offsets,
                 cu_seqlens=attn_metadata.query_start_loc,
                 initial_states=initial_states,
                 return_varlen_states=True,

vllm/model_executor/models/bamba.py

@@ -96,7 +96,6 @@ def __init__(self,
                                 head_dim=config.mamba_d_head,
                                 rms_norm_eps=config.rms_norm_eps,
                                 activation=config.hidden_act,
-                                chunk_size=config.mamba_chunk_size,
                                 quant_config=quant_config)
 
         self.feed_forward = BambaMLP(config, quant_config=quant_config)
@@ -110,7 +109,6 @@ def forward(
         hidden_states: torch.Tensor,
         residual: Optional[torch.Tensor],
         mamba_cache_params: MambaCacheParams,
-        sequence_idx: Optional[torch.Tensor] = None,
         mamba2_metadata: Optional[Mamba2Metadata] = None,
         **kwargs,
     ):
@@ -122,7 +120,7 @@ def forward(
                 hidden_states, residual)
 
         hidden_states = self.mamba(hidden_states, mamba_cache_params,
-                                   sequence_idx, mamba2_metadata)
+                                   mamba2_metadata)
         # Fully Connected
         hidden_states, residual = self.pre_ff_layernorm(
             hidden_states, residual)
@@ -312,33 +310,14 @@ def forward(
         inputs_embeds: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
 
-        # pass a sequence index tensor, that is required for
-        # proper continuous batching computation including
-        # chunked prefill
-        seq_idx = None
-        mamba2_metadata = None
         attn_metadata = get_forward_context().attn_metadata
-        if attn_metadata.num_prefills > 0:
-            seq_idx = torch.zeros_like(input_ids, dtype=torch.int32)
-            for i, (srt, end) in enumerate(
-                    zip(
-                        attn_metadata.query_start_loc,
-                        attn_metadata.query_start_loc[1:],
-                    )):
-                seq_idx[srt:end] = i
-            seq_idx.unsqueeze_(0)
-
-            # compute metadata for chunked prefill.
-            # actually this is only needed if there are
-            # initial states, but this is determinable
-            # only from attention metadata yet
-            # unavailable from the current top-level forward.
-            # Rather than complicating things to extract said
-            # metadata, we simply just compute redundently and
-            # will be silently ignored inside the mamba kernels.
-            # if not needed.
-            mamba2_metadata = prepare_mamba2_metadata(
-                seq_idx, self.config.mamba_chunk_size)
+
+        mamba2_metadata = prepare_mamba2_metadata(
+            chunk_size=self.config.mamba_chunk_size,
+            has_prefills=attn_metadata.num_prefills > 0,
+            input_ids=input_ids,
+            query_start_loc=attn_metadata.query_start_loc,
+        )
 
         if get_pp_group().is_first_rank:
             if inputs_embeds is not None:
@@ -368,7 +347,6 @@ def forward(
                 hidden_states=hidden_states,
                 residual=residual,
                 mamba_cache_params=layer_mamba_cache_params,
-                sequence_idx=seq_idx,
                 mamba2_metadata=mamba2_metadata,
             )