A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[Model] Reduce redundant computations in mamba2 blocks for Bamba-9B (#15423)
Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com> Co-authored-by: Yu Chin Fabian Lim <flim@sg.ibm.com>
Chih-Chieh Yang committed
Apr 10, 2025 at 15:07 UTC
daefed052ccc2811993262cfbf6b746b56021b06
8 files changed
+186
-132
vllm/model_executor/layers/mamba/mamba2_metadata.py
new
+109
@@ -0,0 +1,109 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import math
3
+from dataclasses import dataclass
4
+
5
+import torch
6
+
7
+from vllm.attention.backends.abstract import AttentionMetadata
8
+from vllm.attention.backends.flash_attn import FlashAttentionMetadata
9
+from vllm.attention.backends.placeholder_attn import (
10
+ PlaceholderAttentionMetadata)
11
+from vllm.attention.backends.xformers import XFormersMetadata
12
+
13
+
14
+@dataclass
15
+class Mamba2Metadata:
16
+ has_prefill: bool
17
+
18
+ has_initial_states: torch.Tensor
19
+ prep_initial_states: bool
20
+
21
+ chunk_size: int
22
+ seq_idx: torch.Tensor
23
+ chunk_indices: torch.Tensor
24
+ chunk_offsets: torch.Tensor
25
+
26
+
27
+def _seq_idx_to_chunk_indices_offsets(seq_idx, chunk_size: int):
28
+
29
+ # convert seq_idx to chunk indices and offsets
30
+ # - derive the cu_seqlens
31
+ _, cu_seqlens = torch.where(seq_idx.diff())
32
+ cu_seqlens += 1
33
+
34
+ # outputs will have length expansion of chunks that do not divide
35
+ # chunk_size
36
+ N = math.ceil(seq_idx.shape[-1] / chunk_size) + (cu_seqlens % chunk_size
37
+ > 0).sum()
38
+ chunk_indices = torch.arange(N, dtype=torch.int, device=seq_idx.device)
39
+ chunk_offsets = torch.zeros((N, ), dtype=torch.int, device=seq_idx.device)
40
+
41
+ cu_seqlens = cu_seqlens.tolist() + [seq_idx.shape[-1]]
42
+ p = 0 # num of insertions
43
+ for s, e in zip(cu_seqlens[:-1], cu_seqlens[1:]):
44
+
45
+ # if does not divide chunk_size, then there is one chunk insertion
46
+ p += (s % chunk_size > 0)
47
+
48
+ # get the dimensions
49
+ # - the + 1 for _e is to shift the boundary by one chunk
50
+ # - this shifting is not needed if chunk_size divides e
51
+ _s, _e = s // chunk_size + p, e // chunk_size + p + (e % chunk_size
52
+ > 0)
53
+
54
+ # adjust inidces and offsets
55
+ chunk_indices[_s:_e] -= p
56
+ chunk_offsets[_s] = s % chunk_size
57
+
58
+ return chunk_indices, chunk_offsets
59
+
60
+
61
+def prepare_mamba2_metadata(
62
+ chunk_size: int,
63
+ input_ids: torch.Tensor,
64
+ attn_metadata: AttentionMetadata,
65
+) -> Mamba2Metadata:
66
+
67
+ # Need flags to indicate if there are initial states
68
+ # currently we really only support the FlashAttention backend
69
+ has_initial_states = None
70
+ prep_initial_states = False
71
+ if (isinstance(attn_metadata, (FlashAttentionMetadata, XFormersMetadata,
72
+ PlaceholderAttentionMetadata))
73
+ and attn_metadata.context_lens_tensor is not None):
74
+ has_initial_states = attn_metadata.context_lens_tensor > 0
75
+ # precompute flag to avoid device syncs later in mamba2 forwards
76
+ prep_initial_states = torch.any(has_initial_states).item()
77
+
78
+ has_prefill = attn_metadata.num_prefills > 0
79
+
80
+ seq_idx = None
81
+ chunk_indices, chunk_offsets = None, None
82
+ if has_prefill:
83
+ seq_idx = torch.zeros_like(input_ids, dtype=torch.int32)
84
+ for i, (srt, end) in enumerate(
85
+ zip(
86
+ attn_metadata.query_start_loc,
87
+ attn_metadata.query_start_loc[1:],
88
+ )):
89
+ seq_idx[srt:end] = i
90
+ seq_idx.unsqueeze_(0)
91
+
92
+ # compute metadata for chunked prefill.
93
+ # actually this is only needed if there are initial states,
94
+ # but this is determinable only from attention metadata yet
95
+ # unavailable from the top-level model forward. Rather than
96
+ # complicating things to extract said metadata, we simply just
97
+ # compute them once at the top level model forward and reuse
98
+ # them in mamba layers. If not needed, they will be ignored
99
+ # inside mamba kernels.
100
+ chunk_indices, chunk_offsets = _seq_idx_to_chunk_indices_offsets(
101
+ seq_idx, chunk_size)
102
+
103
+ return Mamba2Metadata(has_prefill=has_prefill,
104
+ has_initial_states=has_initial_states,
105
+ prep_initial_states=prep_initial_states,
106
+ chunk_size=chunk_size,
107
+ seq_idx=seq_idx,
108
+ chunk_indices=chunk_indices,
109
+ chunk_offsets=chunk_offsets)
vllm/model_executor/layers/mamba/mamba_mixer2.py
+20
-32
@@ -6,10 +6,6 @@ import torch
6
from torch import nn
7
8
from vllm.attention.backends.abstract import AttentionMetadata
9
-from vllm.attention.backends.flash_attn import FlashAttentionMetadata
10
-from vllm.attention.backends.placeholder_attn import (
11
- PlaceholderAttentionMetadata)
12
-from vllm.attention.backends.xformers import XFormersMetadata
9
from vllm.distributed import (divide, get_tensor_model_parallel_rank,
10
get_tensor_model_parallel_world_size,
11
tensor_model_parallel_all_gather,
@@ -18,6 +14,7 @@ from vllm.forward_context import get_forward_context
14
from vllm.model_executor.custom_op import CustomOp
15
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
16
RowParallelLinear)
17
+from vllm.model_executor.layers.mamba.mamba2_metadata import Mamba2Metadata
18
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
19
causal_conv1d_fn, causal_conv1d_update)
20
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
@@ -221,7 +218,6 @@ class MambaMixer2(CustomOp):
218
head_dim: int = 64,
219
rms_norm_eps: float = 1e-5,
220
activation="silu",
224
- chunk_size: int = 256,
221
quant_config: Optional[QuantizationConfig] = None):
222
super().__init__()
223
@@ -257,7 +253,6 @@ class MambaMixer2(CustomOp):
253
self.ssm_state_size = ssm_state_size
254
self.activation = activation
255
260
- self.chunk_size = chunk_size
256
self.intermediate_size = intermediate_size
257
self.head_dim = head_dim
258
self.num_heads = num_heads
@@ -388,25 +383,17 @@ class MambaMixer2(CustomOp):
383
self,
384
hidden_states: torch.Tensor,
385
mamba_cache_params: MambaCacheParams,
391
- sequence_idx: Optional[torch.Tensor] = None,
386
+ mamba2_metadata: Mamba2Metadata,
387
):
388
+ # mamba2_metadata contains metadata necessary for the mamba2 triton
389
+ # kernels to operate in continuous batching and in chunked prefill
390
+ # modes; they are computed at top-level model forward since they
391
+ # are the same and reused for all mamba layers in the same iteration
392
attn_metadata: AttentionMetadata = get_forward_context().attn_metadata
393
394
seq_len, _ = hidden_states.shape
395
groups_time_state_size = self.n_groups * self.ssm_state_size
396
398
- # detect if there are prefills
399
- has_prefill = attn_metadata.num_prefills > 0
400
-
401
- # - also need flags to indicate if there are initial states
402
- # - currently we really only support the FlashAttention backend
403
- has_initial_states = None
404
- if (isinstance(attn_metadata,
405
- (FlashAttentionMetadata, XFormersMetadata,
406
- PlaceholderAttentionMetadata))
407
- and attn_metadata.context_lens_tensor is not None):
408
- has_initial_states = attn_metadata.context_lens_tensor > 0
409
-
397
# 1. Gated MLP's linear projection
398
projected_states, _ = self.in_proj(hidden_states)
399
gate, hidden_states_B_C, dt = torch.split(
@@ -423,7 +410,7 @@ class MambaMixer2(CustomOp):
410
conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),
411
self.conv1d.weight.size(2))
412
426
- if has_prefill:
413
+ if mamba2_metadata.has_prefill:
414
# |---------- N-1 iteration --------|
415
# |---------------- N iteration ---------------------|
416
# |- tokenA -|......................|-- newTokens ---|
@@ -439,7 +426,7 @@ class MambaMixer2(CustomOp):
426
self.conv1d.bias,
427
activation=self.activation,
428
conv_states=mamba_cache_params.conv_state,
442
- has_initial_state=has_initial_states,
429
+ has_initial_state=mamba2_metadata.has_initial_states,
430
cache_indices=mamba_cache_params.state_indices_tensor,
431
query_start_loc=attn_metadata.query_start_loc).transpose(
432
0, 1)[:seq_len]
@@ -467,16 +454,15 @@ class MambaMixer2(CustomOp):
454
)
455
456
# 3. State Space Model sequence transformation
470
- if has_prefill:
471
-
457
+ if mamba2_metadata.has_prefill:
458
initial_states = None
473
- if has_initial_states is not None and torch.any(
474
- has_initial_states):
475
- zero_init_indices = mamba_cache_params.state_indices_tensor[
476
- ~has_initial_states]
477
- mamba_cache_params.ssm_state[zero_init_indices] = 0
478
- initial_states = mamba_cache_params.ssm_state[
479
- mamba_cache_params.state_indices_tensor]
459
+ if (mamba2_metadata.has_initial_states is not None
460
+ and mamba2_metadata.prep_initial_states):
461
+ # making a copy of the states
462
+ initial_states = torch.where(
463
+ mamba2_metadata.has_initial_states[:, None, None, None],
464
+ mamba_cache_params.ssm_state[
465
+ mamba_cache_params.state_indices_tensor], 0)
466
467
scan_output, varlen_state = mamba_chunk_scan_combined(
468
hidden_states.view(1, seq_len, self.num_heads // self.tp_size,
@@ -485,11 +471,13 @@ class MambaMixer2(CustomOp):
471
self.A,
472
B.view(1, seq_len, self.n_groups // self.tp_size, -1),
473
C.view(1, seq_len, self.n_groups // self.tp_size, -1),
488
- chunk_size=self.chunk_size,
474
+ chunk_size=mamba2_metadata.chunk_size,
475
D=self.D,
476
z=None,
477
dt_bias=self.dt_bias,
492
- seq_idx=sequence_idx,
478
+ seq_idx=mamba2_metadata.seq_idx,
479
+ chunk_indices=mamba2_metadata.chunk_indices,
480
+ chunk_offsets=mamba2_metadata.chunk_offsets,
481
cu_seqlens=attn_metadata.query_start_loc,
482
initial_states=initial_states,
483
return_varlen_states=True,
vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py
+10
-41
@@ -5,8 +5,6 @@
5
6
# ruff: noqa: E501,SIM102
7
8
-import math
9
-
8
import torch
9
import triton
10
import triton.language as tl
@@ -442,40 +440,6 @@ def _chunk_scan_fwd_kernel(
440
(offs_out_n[None, :] < hdim))
441
442
445
-def _seq_idx_to_chunk_indices_offsets(seq_idx, chunk_size: int):
446
-
447
- # convert seq_idx to chunk indices and offsets
448
- # - derive the cu_seqlens
449
- _, cu_seqlens = torch.where(seq_idx.diff())
450
- cu_seqlens += 1
451
-
452
- # outputs will have length expansion of chunks that do not divide
453
- # chunk_size
454
- N = math.ceil(seq_idx.shape[-1] / chunk_size) + (cu_seqlens % chunk_size
455
- > 0).sum()
456
- chunk_indices = torch.arange(N, dtype=torch.int, device=seq_idx.device)
457
- chunk_offsets = torch.zeros((N, ), dtype=torch.int, device=seq_idx.device)
458
-
459
- cu_seqlens = cu_seqlens.tolist() + [seq_idx.shape[-1]]
460
- p = 0 # num of insertions
461
- for s, e in zip(cu_seqlens[:-1], cu_seqlens[1:]):
462
-
463
- # if does not divide chunk_size, then there is one chunk insertion
464
- p += (s % chunk_size > 0)
465
-
466
- # get the dimensions
467
- # - the + 1 for _e is to shift the boundary by one chunk
468
- # - this shifting is not needed if chunk_size divides e
469
- _s, _e = s // chunk_size + p, e // chunk_size + p + (e % chunk_size
470
- > 0)
471
-
472
- # adjust inidces and offsets
473
- chunk_indices[_s:_e] -= p
474
- chunk_offsets[_s] = s % chunk_size
475
-
476
- return chunk_indices, chunk_offsets
477
-
478
-
443
def _chunk_scan_fwd(
444
cb,
445
x,
@@ -486,6 +450,8 @@ def _chunk_scan_fwd(
450
D=None,
451
z=None,
452
seq_idx=None,
453
+ chunk_indices=None,
454
+ chunk_offsets=None,
455
initial_states=None,
456
):
457
batch, seqlen, nheads, headdim = x.shape
@@ -502,7 +468,6 @@ def _chunk_scan_fwd(
468
assert dA_cumsum.shape == (batch, nheads, nchunks, chunk_size)
469
assert states.shape == (batch, nchunks, nheads, headdim, dstate)
470
505
- chunk_indices, chunk_offsets = None, None
471
if seq_idx is not None:
472
assert seq_idx.shape == (batch, seqlen)
473
@@ -510,15 +475,19 @@ def _chunk_scan_fwd(
475
# with initial states, we need to take care of how
476
# seq_idx crosses the boundaries
477
assert batch == 1, "chunk scan only supports initial states with batch 1"
513
- assert initial_states.shape == (seq_idx[0].max() + 1, nheads,
514
- headdim, dstate)
478
479
if initial_states.shape[0] == 1:
480
# no in this case no point to use initial states
481
initial_states = None
482
else:
520
- chunk_indices, chunk_offsets = _seq_idx_to_chunk_indices_offsets(
521
- seq_idx, chunk_size)
483
+ assert chunk_indices is not None and chunk_offsets is not None, \
484
+ (
485
+ "chunk_indices and chunk_offsets should have been set"
486
+ )
487
+ else:
488
+ chunk_indices, chunk_offsets = None, None
489
+ else:
490
+ chunk_indices, chunk_offsets = None, None
491
492
# Allocates output.
493
out = torch.empty(batch,
vllm/model_executor/layers/mamba/ops/ssd_combined.py
+9
-1
@@ -30,6 +30,8 @@ def _mamba_chunk_scan_combined_fwd(x,
30
dt_bias=None,
31
initial_states=None,
32
seq_idx=None,
33
+ chunk_indices=None,
34
+ chunk_offsets=None,
35
cu_seqlens=None,
36
dt_softplus=False,
37
dt_limit=(0.0, float("inf"))):
@@ -96,7 +98,7 @@ def _mamba_chunk_scan_combined_fwd(x,
98
# 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
99
# (middle term of factorization of off-diag blocks; A terms)
100
# - for handling chunked prefill, this requires i) initial_states
99
- # ii) seq_idx and iii) has_cu_seqlens to be all specified.
101
+ # ii) seq_idx and iii) is_cont_batched to be all specified.
102
# - When a new seq_idx is detected, we will stop passing the prev_state
103
# and switch accordingly to the init_state corresponding to the new seq_idx.
104
# - this will ensure that states will be updated with the rightmost flushed seq_idx
@@ -141,6 +143,8 @@ def _mamba_chunk_scan_combined_fwd(x,
143
D=D,
144
z=z,
145
seq_idx=seq_idx,
146
+ chunk_indices=chunk_indices,
147
+ chunk_offsets=chunk_offsets,
148
initial_states=initial_states,
149
)
150
if cu_seqlens is None:
@@ -170,6 +174,8 @@ def mamba_chunk_scan_combined(x,
174
dt_bias=None,
175
initial_states=None,
176
seq_idx=None,
177
+ chunk_indices=None,
178
+ chunk_offsets=None,
179
cu_seqlens=None,
180
dt_softplus=False,
181
dt_limit=(0.0, float("inf")),
@@ -210,6 +216,8 @@ def mamba_chunk_scan_combined(x,
216
dt_bias=dt_bias,
217
initial_states=initial_states,
218
seq_idx=seq_idx,
219
+ chunk_indices=chunk_indices,
220
+ chunk_offsets=chunk_offsets,
221
cu_seqlens=cu_seqlens,
222
dt_softplus=dt_softplus,
223
dt_limit=dt_limit)
vllm/model_executor/layers/mamba/ops/ssd_state_passing.py
-2
@@ -150,8 +150,6 @@ def _state_passing_fwd(
150
# are used for continuous batching. In which case we
151
# require seq_idx to be provided
152
assert seq_idx is not None, ""
153
- assert initial_states.shape == (seq_idx.max().item() + 1, nheads,
154
- dim)
153
else:
154
# - this is the regular batching case, where initial
155
# states are used are for each example of the batch.
vllm/model_executor/models/bamba.py
+13
-19
@@ -18,6 +18,8 @@ from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
18
QKVParallelLinear,
19
RowParallelLinear)
20
from vllm.model_executor.layers.logits_processor import LogitsProcessor
21
+from vllm.model_executor.layers.mamba.mamba2_metadata import (
22
+ Mamba2Metadata, prepare_mamba2_metadata)
23
from vllm.model_executor.layers.mamba.mamba_mixer2 import (
24
MambaMixer2, extra_groups_for_head_shards)
25
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -94,7 +96,6 @@ class BambaMixerDecoderLayer(nn.Module):
96
head_dim=config.mamba_d_head,
97
rms_norm_eps=config.rms_norm_eps,
98
activation=config.hidden_act,
97
- chunk_size=config.mamba_chunk_size,
99
quant_config=quant_config)
100
101
self.feed_forward = BambaMLP(config, quant_config=quant_config)
@@ -108,7 +109,7 @@ class BambaMixerDecoderLayer(nn.Module):
109
hidden_states: torch.Tensor,
110
residual: Optional[torch.Tensor],
111
mamba_cache_params: MambaCacheParams,
111
- sequence_idx: Optional[torch.Tensor] = None,
112
+ mamba2_metadata: Mamba2Metadata,
113
**kwargs,
114
):
115
if residual is None:
@@ -119,7 +120,7 @@ class BambaMixerDecoderLayer(nn.Module):
120
hidden_states, residual)
121
122
hidden_states = self.mamba(hidden_states, mamba_cache_params,
122
- sequence_idx)
123
+ mamba2_metadata)
124
# Fully Connected
125
hidden_states, residual = self.pre_ff_layernorm(
126
hidden_states, residual)
@@ -259,7 +260,7 @@ class BambaModel(nn.Module):
260
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
261
super().__init__()
262
262
- config = vllm_config.model_config.hf_config
263
+ config: BambaConfig = vllm_config.model_config.hf_config
264
cache_config = vllm_config.cache_config
265
quant_config = vllm_config.quant_config
266
lora_config = vllm_config.lora_config
@@ -309,20 +310,13 @@ class BambaModel(nn.Module):
310
inputs_embeds: Optional[torch.Tensor] = None,
311
) -> torch.Tensor:
312
312
- # pass a sequence index tensor, that is required for
313
- # proper continuous batching computation including
314
- # chunked prefill
315
- seq_idx = None
313
attn_metadata = get_forward_context().attn_metadata
317
- if attn_metadata.num_prefills > 0:
318
- seq_idx = torch.zeros_like(input_ids, dtype=torch.int32)
319
- for i, (srt, end) in enumerate(
320
- zip(
321
- attn_metadata.query_start_loc,
322
- attn_metadata.query_start_loc[1:],
323
- )):
324
- seq_idx[srt:end] = i
325
- seq_idx.unsqueeze_(0)
314
+
315
+ mamba2_metadata = prepare_mamba2_metadata(
316
+ chunk_size=self.config.mamba_chunk_size,
317
+ input_ids=input_ids,
318
+ attn_metadata=attn_metadata,
319
+ )
320
321
if get_pp_group().is_first_rank:
322
if inputs_embeds is not None:
@@ -352,7 +346,7 @@ class BambaModel(nn.Module):
346
hidden_states=hidden_states,
347
residual=residual,
348
mamba_cache_params=layer_mamba_cache_params,
355
- sequence_idx=seq_idx,
349
+ mamba2_metadata=mamba2_metadata,
350
)
351
352
if not get_pp_group().is_last_rank:
@@ -555,4 +549,4 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
549
def load_weights(self, weights: Iterable[Tuple[str,
550
torch.Tensor]]) -> Set[str]:
551
loader = AutoWeightsLoader(self)
558
- return loader.load_weights(weights)
\ No newline at end of file
552
+ return loader.load_weights(weights)
vllm/model_executor/models/mamba2.py
+11
-17
@@ -13,6 +13,8 @@ from vllm.distributed.parallel_state import get_pp_group
13
from vllm.forward_context import get_forward_context
14
from vllm.model_executor.layers.layernorm import RMSNorm
15
from vllm.model_executor.layers.logits_processor import LogitsProcessor
16
+from vllm.model_executor.layers.mamba.mamba2_metadata import (
17
+ Mamba2Metadata, prepare_mamba2_metadata)
18
from vllm.model_executor.layers.mamba.mamba_mixer2 import (
19
MambaMixer2, extra_groups_for_head_shards)
20
from vllm.model_executor.layers.quantization.base_config import (
@@ -57,7 +59,6 @@ class Mamba2DecoderLayer(nn.Module):
59
head_dim=config.head_dim,
60
rms_norm_eps=config.layer_norm_epsilon,
61
activation=config.hidden_act,
60
- chunk_size=config.chunk_size,
62
quant_config=quant_config)
63
64
self.norm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
@@ -67,7 +68,7 @@ class Mamba2DecoderLayer(nn.Module):
68
hidden_states: torch.Tensor,
69
residual: Optional[torch.Tensor],
70
mamba_cache_params: MambaCacheParams,
70
- sequence_idx: Optional[torch.Tensor],
71
+ mamba2_metadata: Mamba2Metadata,
72
**kwargs,
73
):
74
if residual is None:
@@ -77,7 +78,7 @@ class Mamba2DecoderLayer(nn.Module):
78
hidden_states, residual = self.norm(hidden_states, residual)
79
80
hidden_states = self.mixer(hidden_states, mamba_cache_params,
80
- sequence_idx)
81
+ mamba2_metadata)
82
return hidden_states, residual
83
84
@@ -138,20 +139,13 @@ class Mamba2Model(nn.Module):
139
hidden_states = intermediate_tensors["hidden_states"]
140
residual = intermediate_tensors["residual"]
141
141
- # pass a sequence index tensor, that is required for
142
- # proper continuous batching computation including
143
- # chunked prefill
144
- seq_idx = None
142
attn_metadata: AttentionMetadata = get_forward_context().attn_metadata
146
- if attn_metadata.num_prefills > 0:
147
- seq_idx = torch.zeros_like(input_ids, dtype=torch.int32)
148
- for i, (srt, end) in enumerate(
149
- zip(
150
- attn_metadata.query_start_loc,
151
- attn_metadata.query_start_loc[1:],
152
- )):
153
- seq_idx[srt:end] = i
154
- seq_idx.unsqueeze_(0)
143
+
144
+ mamba2_metadata = prepare_mamba2_metadata(
145
+ chunk_size=self.config.chunk_size,
146
+ input_ids=input_ids,
147
+ attn_metadata=attn_metadata,
148
+ )
149
150
for i in range(len(self.layers)):
151
layer = self.layers[i]
@@ -162,7 +156,7 @@ class Mamba2Model(nn.Module):
156
residual=residual,
157
mamba_cache_params=mamba_cache_params.at_layer_idx(
158
i - self.start_layer),
165
- sequence_idx=seq_idx)
159
+ mamba2_metadata=mamba2_metadata)
160
161
if not get_pp_group().is_last_rank:
162
return IntermediateTensors({
vllm/model_executor/models/zamba2.py
+14
-20
@@ -25,6 +25,8 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
25
ReplicatedLinear,
26
RowParallelLinear)
27
from vllm.model_executor.layers.logits_processor import LogitsProcessor
28
+from vllm.model_executor.layers.mamba.mamba2_metadata import (
29
+ Mamba2Metadata, prepare_mamba2_metadata)
30
from vllm.model_executor.layers.mamba.mamba_mixer2 import (
31
MambaMixer2, extra_groups_for_head_shards)
32
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -495,7 +497,6 @@ class Zamba2MambaDecoderLayer(nn.Module):
497
head_dim=intermediate_size // config.n_mamba_heads,
498
rms_norm_eps=config.rms_norm_eps,
499
activation="silu",
498
- chunk_size=config.chunk_size,
500
quant_config=quant_config,
501
)
502
@@ -507,7 +508,7 @@ class Zamba2MambaDecoderLayer(nn.Module):
508
self,
509
hidden_states: torch.Tensor,
510
mamba_cache_params: MambaCacheParams,
510
- sequence_idx: Optional[torch.Tensor] = None,
511
+ mamba2_metadata: Mamba2Metadata,
512
transformer_hidden_states: Optional[torch.Tensor] = None,
513
positions: Optional[torch.Tensor] = None,
514
original_hidden_states: Optional[torch.Tensor] = None,
@@ -547,7 +548,7 @@ class Zamba2MambaDecoderLayer(nn.Module):
548
hidden_states = self.mamba(
549
hidden_states,
550
mamba_cache_params=mamba_cache_params,
550
- sequence_idx=sequence_idx,
551
+ mamba2_metadata=mamba2_metadata,
552
)
553
554
# residual connection after mamba
@@ -594,8 +595,8 @@ class Zamba2HybridLayer(nn.Module):
595
hidden_states: torch.Tensor,
596
original_hidden_states: torch.Tensor,
597
positions: torch.Tensor,
597
- mamba_cache_params: Optional[MambaCacheParams] = None,
598
- sequence_idx: Optional[torch.Tensor] = None,
598
+ mamba_cache_params: MambaCacheParams,
599
+ mamba2_metadata: Mamba2Metadata,
600
) -> torch.Tensor:
601
"""Forward pass through the hybrid layer.
602
@@ -634,7 +635,7 @@ class Zamba2HybridLayer(nn.Module):
635
hidden_states,
636
transformer_hidden_states=transformer_hidden_states,
637
mamba_cache_params=mamba_cache_params,
637
- sequence_idx=sequence_idx,
638
+ mamba2_metadata=mamba2_metadata,
639
)
640
641
return layer_outputs
@@ -747,20 +748,13 @@ class Zamba2Model(nn.Module):
748
inputs_embeds = self.get_input_embeddings(input_ids)
749
hidden_states = inputs_embeds
750
750
- # pass a sequence index tensor, that is required for
751
- # proper continuous batching computation including
752
- # chunked prefill
753
- seq_idx = None
751
attn_metadata = get_forward_context().attn_metadata
755
- if attn_metadata.num_prefills > 0:
756
- seq_idx = torch.zeros_like(input_ids, dtype=torch.int32)
757
- for i, (srt, end) in enumerate(
758
- zip(
759
- attn_metadata.query_start_loc,
760
- attn_metadata.query_start_loc[1:],
761
- )):
762
- seq_idx[srt:end] = i
763
- seq_idx.unsqueeze_(0)
752
+
753
+ mamba2_metadata = prepare_mamba2_metadata(
754
+ chunk_size=self.config.chunk_size,
755
+ input_ids=input_ids,
756
+ attn_metadata=attn_metadata,
757
+ )
758
759
# Process through layers
760
original_hidden_states = torch.clone(hidden_states)
@@ -770,7 +764,7 @@ class Zamba2Model(nn.Module):
764
original_hidden_states=original_hidden_states,
765
positions=positions,
766
mamba_cache_params=mamba_cache_params.at_layer_idx(layer_idx),
773
- sequence_idx=seq_idx,
767
+ mamba2_metadata=mamba2_metadata,
768
)
769
hidden_states = layer_outputs
770