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 14, 05:30 UTC
[V1] Remove V0 code paths for Hybrid models (#25400)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Sep 23, 2025 at 17:26 UTC
a903669e10cc98507fa5c2ae099b7161f7140cf7
31 files changed
+359
-2303
tests/models/language/generation/test_hybrid.py
+19
-36
@@ -20,7 +20,9 @@ pytestmark = pytest.mark.hybrid_model
20
SSM_MODELS = [
21
"state-spaces/mamba-130m-hf",
22
"tiiuae/falcon-mamba-tiny-dev",
23
- "yujiepan/mamba2-codestral-v0.1-tiny-random",
23
+ # mamba2-codestral in transformers is broken pending:
24
+ # https://github.com/huggingface/transformers/pull/40861
25
+ #"yujiepan/mamba2-codestral-v0.1-tiny-random",
26
]
27
28
HYBRID_MODELS = [
@@ -31,18 +33,7 @@ HYBRID_MODELS = [
33
"ibm-granite/granite-4.0-tiny-preview",
34
"tiiuae/Falcon-H1-0.5B-Base",
35
"LiquidAI/LFM2-1.2B",
34
-]
35
-
36
-V1_SUPPORTED_MODELS = [
37
- "state-spaces/mamba-130m-hf",
38
- "ai21labs/Jamba-tiny-dev",
39
- "pfnet/plamo-2-1b",
40
- "yujiepan/mamba2-codestral-v0.1-tiny-random",
41
- "Zyphra/Zamba2-1.2B-instruct",
42
- "hmellor/tiny-random-BambaForCausalLM",
43
- "ibm-granite/granite-4.0-tiny-preview",
44
- "tiiuae/Falcon-H1-0.5B-Base",
45
- "LiquidAI/LFM2-1.2B",
36
+ "tiny-random/qwen3-next-moe",
37
]
38
39
FULL_CUDA_GRAPH_MODELS = [
@@ -51,10 +42,6 @@ FULL_CUDA_GRAPH_MODELS = [
42
"Zyphra/Zamba2-1.2B-instruct",
43
]
44
54
-V0_UNSUPPORTED_MODELS = [
55
- "LiquidAI/LFM2-1.2B",
56
-]
57
-
45
FP32_STATE_MODELS = [
46
"state-spaces/mamba-130m-hf",
47
"Zyphra/Zamba2-1.2B-instruct",
@@ -88,20 +75,16 @@ def test_models(
75
hf_outputs = hf_model.generate_greedy_logprobs_limit(
76
example_prompts, max_tokens, num_logprobs)
77
91
- if model in V1_SUPPORTED_MODELS:
92
- with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
93
- vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
94
- example_prompts, max_tokens, num_logprobs)
95
- else:
96
- vllm_v1_outputs = None
78
+ with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
79
+ vllm_outputs = vllm_model.generate_greedy_logprobs(
80
+ example_prompts, max_tokens, num_logprobs)
81
98
- if model in V1_SUPPORTED_MODELS:
99
- check_logprobs_close(
100
- outputs_0_lst=hf_outputs,
101
- outputs_1_lst=vllm_v1_outputs,
102
- name_0="hf",
103
- name_1="vllm-v1",
104
- )
82
+ check_logprobs_close(
83
+ outputs_0_lst=hf_outputs,
84
+ outputs_1_lst=vllm_outputs,
85
+ name_0="hf",
86
+ name_1="vllm",
87
+ )
88
89
90
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
@@ -299,14 +282,14 @@ def test_full_cuda_graph(
282
example_prompts, max_tokens, num_logprobs)
283
284
with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
302
- vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
285
+ vllm_outputs = vllm_model.generate_greedy_logprobs(
286
example_prompts, max_tokens, num_logprobs)
287
288
check_logprobs_close(
289
outputs_0_lst=hf_outputs,
307
- outputs_1_lst=vllm_v1_outputs,
290
+ outputs_1_lst=vllm_outputs,
291
name_0="hf",
309
- name_1="vllm-v1",
292
+ name_1="vllm",
293
)
294
295
@@ -340,12 +323,12 @@ def test_fp32_cache_state(
323
with vllm_runner(model,
324
max_num_seqs=MAX_NUM_SEQS,
325
**{cache_dtype_param: "float32"}) as vllm_model:
343
- vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
326
+ vllm_outputs = vllm_model.generate_greedy_logprobs(
327
example_prompts, max_tokens, num_logprobs)
328
329
check_logprobs_close(
330
outputs_0_lst=hf_outputs,
348
- outputs_1_lst=vllm_v1_outputs,
331
+ outputs_1_lst=vllm_outputs,
332
name_0="hf",
350
- name_1="vllm-v1",
333
+ name_1="vllm",
334
)
tests/models/registry.py
+6
-7
@@ -312,14 +312,12 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
312
"PersimmonForCausalLM": _HfExamplesInfo("adept/persimmon-8b-chat"),
313
"PhiForCausalLM": _HfExamplesInfo("microsoft/phi-2"),
314
"Phi3ForCausalLM": _HfExamplesInfo("microsoft/Phi-3-mini-4k-instruct"),
315
- "Phi4FlashForCausalLM": _HfExamplesInfo("microsoft/Phi-4-mini-flash-reasoning", # noqa: E501
316
- trust_remote_code=True,
317
- v0_only=True,
318
- max_model_len=10240),
315
"PhiMoEForCausalLM": _HfExamplesInfo("microsoft/Phi-3.5-MoE-instruct",
316
trust_remote_code=True),
317
"Plamo2ForCausalLM": _HfExamplesInfo("pfnet/plamo-2-1b",
322
- trust_remote_code=True),
318
+ max_transformers_version="4.55.4",
319
+ transformers_version_reason="HF model uses remote code that is not compatible with latest Transformers", # noqa: E501
320
+ trust_remote_code=True),
321
"QWenLMHeadModel": _HfExamplesInfo("Qwen/Qwen-7B-Chat",
322
max_transformers_version="4.53",
323
transformers_version_reason="HF model uses remote code that is not compatible with latest Transformers", # noqa: E501
@@ -330,7 +328,8 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
328
"Qwen3ForCausalLM": _HfExamplesInfo("Qwen/Qwen3-8B"),
329
"Qwen3MoeForCausalLM": _HfExamplesInfo("Qwen/Qwen3-30B-A3B"),
330
"Qwen3NextForCausalLM": _HfExamplesInfo("Qwen/Qwen3-Next-80B-A3B-Instruct",
333
- min_transformers_version="4.56.2"),
331
+ extras={"tiny-random": "tiny-random/qwen3-next-moe"}, # noqa: E501
332
+ min_transformers_version="4.56.3"),
333
"RWForCausalLM": _HfExamplesInfo("tiiuae/falcon-40b"),
334
"SeedOssForCausalLM": _HfExamplesInfo("ByteDance-Seed/Seed-OSS-36B-Instruct", # noqa: E501
335
trust_remote_code=True,
@@ -644,7 +643,7 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
643
trust_remote_code=True,
644
speculative_model="XiaomiMiMo/MiMo-7B-RL"),
645
"Qwen3NextMTP": _HfExamplesInfo("Qwen/Qwen3-Next-80B-A3B-Instruct",
647
- min_transformers_version="4.56.2"),
646
+ min_transformers_version="4.56.3"),
647
}
648
649
_TRANSFORMERS_BACKEND_MODELS = {
vllm/model_executor/layers/mamba/abstract.py
+1
-4
@@ -20,10 +20,7 @@ class MambaBase(AttentionLayerBase):
20
21
# Contains the KV cache (mamba state) for the layer
22
# in the shape specified by `self.get_state_shape`.
23
- # The outer list is for v0 PP virtual engine. Though this code path
24
- # only runs for v1, we have to do this to unify with the interface
25
- # of Attention + v0 PP.
26
- kv_cache: list[Iterable[torch.Tensor]]
23
+ kv_cache: tuple[torch.Tensor, ...]
24
25
@abstractmethod
26
def get_state_shape(self) -> Iterable[tuple[int, ...]]:
vllm/model_executor/layers/mamba/linear_attn.py
+40
-67
@@ -15,7 +15,6 @@ import torch.nn.functional as F
15
from einops import rearrange
16
from torch import nn
17
18
-from vllm import envs
18
from vllm.attention import AttentionMetadata
19
from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
20
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
@@ -42,8 +41,6 @@ if TYPE_CHECKING:
41
import torch
42
import torch.distributed
43
45
-from vllm.model_executor.models.minimax_cache import MinimaxCacheParams
46
-
44
45
class MiniMaxText01RMSNormTP(CustomOp):
46
name = "MiniMaxText01RMSNormTP"
@@ -225,11 +222,10 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
222
self.tp_heads:(self.tp_rank + 1) *
223
self.tp_heads].contiguous()
224
228
- if envs.VLLM_USE_V1:
229
- compilation_config = get_current_vllm_config().compilation_config
230
- if prefix in compilation_config.static_forward_context:
231
- raise ValueError(f"Duplicate layer name: {prefix}")
232
- compilation_config.static_forward_context[prefix] = self
225
+ compilation_config = get_current_vllm_config().compilation_config
226
+ if prefix in compilation_config.static_forward_context:
227
+ raise ValueError(f"Duplicate layer name: {prefix}")
228
+ compilation_config.static_forward_context[prefix] = self
229
230
@staticmethod
231
def weight_direct_load(param: torch.Tensor,
@@ -268,8 +264,7 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
264
break
265
if _prefill_idx >= len(state_indices_tensor):
266
break
271
- # prefills are packed at end of batch in V1
272
- offset = attn_metadata.num_decode_tokens if envs.VLLM_USE_V1 else 0
267
+ offset = attn_metadata.num_decode_tokens
268
_start = attn_metadata.query_start_loc[offset + _prefill_idx]
269
_end = attn_metadata.query_start_loc[offset + _prefill_idx + 1]
270
slot_id = state_indices_tensor[offset + _prefill_idx]
@@ -291,10 +286,7 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
286
hidden_decode = self._decode_infer(q, k, v, kv_cache,
287
state_indices_tensor,
288
attn_metadata)
294
- if envs.VLLM_USE_V1:
295
- hidden.insert(0, hidden_decode)
296
- else:
297
- hidden.append(hidden_decode)
289
+ hidden.insert(0, hidden_decode)
290
291
if not hidden:
292
return torch.empty((0, q.size(-1)), device=q.device, dtype=q.dtype)
@@ -304,40 +296,28 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
296
297
def _decode_infer(self, q, k, v, kv_cache, state_indices_tensor,
298
attn_metadata):
307
- if not envs.VLLM_USE_V1:
308
- q = q[attn_metadata.num_prefill_tokens:].unsqueeze(2).contiguous()
309
- k = k[attn_metadata.num_prefill_tokens:].unsqueeze(2).contiguous()
310
- v = v[attn_metadata.num_prefill_tokens:].unsqueeze(2).contiguous()
311
- num_prefills = getattr(attn_metadata, "num_prefills", 0)
312
- slot_id = state_indices_tensor[num_prefills:]
313
- else:
314
- q = q[:attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()
315
- k = k[:attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()
316
- v = v[:attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()
317
- slot_id = state_indices_tensor[:attn_metadata.num_decodes]
299
+ q = q[:attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()
300
+ k = k[:attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()
301
+ v = v[:attn_metadata.num_decode_tokens].unsqueeze(2).contiguous()
302
+ slot_id = state_indices_tensor[:attn_metadata.num_decodes]
303
hidden = linear_decode_forward_triton(q, k, v, kv_cache, self.tp_slope,
304
slot_id, 32)
305
return hidden
306
307
def forward(self, hidden_states: torch.Tensor, output: torch.Tensor,
323
- positions: torch.Tensor,
324
- kv_caches: MinimaxCacheParams) -> None:
325
- if not envs.VLLM_USE_V1:
326
- self._forward(hidden_states, output, positions, kv_caches)
327
- else:
328
- torch.ops.vllm.linear_attention(
329
- hidden_states,
330
- output,
331
- positions,
332
- self.prefix,
333
- )
308
+ positions: torch.Tensor) -> None:
309
+ torch.ops.vllm.linear_attention(
310
+ hidden_states,
311
+ output,
312
+ positions,
313
+ self.prefix,
314
+ )
315
316
def _forward(self, hidden_states: torch.Tensor, output: torch.Tensor,
336
- positions: torch.Tensor,
337
- kv_caches: Optional[MinimaxCacheParams]) -> None:
317
+ positions: torch.Tensor) -> None:
318
forward_context = get_forward_context()
319
attn_metadata: AttentionMetadata = forward_context.attn_metadata
340
- if envs.VLLM_USE_V1 and attn_metadata is not None:
320
+ if attn_metadata is not None:
321
assert isinstance(attn_metadata, dict)
322
attn_metadata = attn_metadata[self.prefix]
323
assert isinstance(attn_metadata, LinearAttentionMetadata)
@@ -351,32 +331,26 @@ class MiniMaxText01LinearAttention(nn.Module, MambaBase):
331
qkvact = torch.nn.functional.silu(qkv32)
332
qkvact = qkvact.view((qkv.shape[0], self.tp_heads, -1))
333
q, k, v = torch.split(qkvact, [self.head_dim] * 3, dim=-1)
354
- if envs.VLLM_USE_V1:
355
- if attn_metadata is not None:
356
- kv_cache = self.kv_cache[forward_context.virtual_engine][0]
357
- state_indices_tensor = attn_metadata.state_indices_tensor
358
-
359
- num_prefills = getattr(attn_metadata, "num_prefills", 0)
360
- if num_prefills > 0:
361
- num_decode_tokens = getattr(attn_metadata,
362
- "num_decode_tokens", 0)
363
- for prefill_idx in range(num_prefills):
364
- q_start = attn_metadata.query_start_loc[
365
- num_decode_tokens + prefill_idx]
366
- q_end = attn_metadata.query_start_loc[num_decode_tokens
367
- + prefill_idx +
368
- 1]
369
- query_len = q_end - q_start
370
- context_len = attn_metadata.seq_lens[
371
- num_decode_tokens + prefill_idx] - query_len
372
- if context_len == 0:
373
- block_to_clear = state_indices_tensor[
374
- num_decode_tokens + prefill_idx]
375
- kv_cache[block_to_clear, ...] = 0
376
- else:
377
- assert kv_caches is not None
378
- kv_cache = kv_caches.minimax_cache
379
- state_indices_tensor = kv_caches.state_indices_tensor
334
+ if attn_metadata is not None:
335
+ kv_cache = self.kv_cache[forward_context.virtual_engine][0]
336
+ state_indices_tensor = attn_metadata.state_indices_tensor
337
+
338
+ num_prefills = getattr(attn_metadata, "num_prefills", 0)
339
+ if num_prefills > 0:
340
+ num_decode_tokens = getattr(attn_metadata, "num_decode_tokens",
341
+ 0)
342
+ for prefill_idx in range(num_prefills):
343
+ q_start = attn_metadata.query_start_loc[num_decode_tokens +
344
+ prefill_idx]
345
+ q_end = attn_metadata.query_start_loc[num_decode_tokens +
346
+ prefill_idx + 1]
347
+ query_len = q_end - q_start
348
+ context_len = attn_metadata.seq_lens[
349
+ num_decode_tokens + prefill_idx] - query_len
350
+ if context_len == 0:
351
+ block_to_clear = state_indices_tensor[num_decode_tokens
352
+ + prefill_idx]
353
+ kv_cache[block_to_clear, ...] = 0
354
355
decode_only = getattr(attn_metadata, "num_prefills", 0) == 0
356
if attn_metadata is None:
@@ -410,8 +384,7 @@ def linear_attention(
384
self = forward_context.no_compile_layers[layer_name]
385
self._forward(hidden_states=hidden_states,
386
output=output,
413
- positions=positions,
414
- kv_caches=None)
387
+ positions=positions)
388
389
390
def linear_attention_fake(
vllm/model_executor/layers/mamba/mamba2_metadata.py
deleted
-177
@@ -1,177 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from dataclasses import dataclass
4
-from typing import Optional, Union
5
-
6
-import numpy as np
7
-import torch
8
-
9
-from vllm.attention.backends.abstract import AttentionMetadata
10
-from vllm.attention.backends.placeholder_attn import (
11
- PlaceholderAttentionMetadata)
12
-from vllm.attention.backends.utils import PAD_SLOT_ID
13
-from vllm.platforms import current_platform
14
-from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
15
-from vllm.v1.attention.backends.mamba2_attn import (
16
- Mamba2AttentionMetadata, _query_start_loc_to_chunk_indices_offsets)
17
-
18
-
19
-@dataclass
20
-class Mamba2Metadata:
21
- prep_initial_states: bool
22
- chunk_size: int
23
-
24
- has_initial_states_p: torch.Tensor
25
- seq_idx_p: torch.Tensor
26
- chunk_indices_p: torch.Tensor
27
- chunk_offsets_p: torch.Tensor
28
- """
29
- With continuous batching layout of `x` in vLLM, to enable a Triton program
30
- to handle a request in parallel, two supporting tensors are used
31
- (batch_ptr, token_chunk_offset_ptr)
32
- BLOCK_M = the # tokens to be handled by a Triton program
33
- (can be customized for different hardware)
34
-
35
- nums_dict:
36
- tracks the data associated with a given value of BLOCK_M
37
- BLOCK_M = #tokens handled by a Triton program
38
- cu_seqlen: total tokens per batch
39
- (used as flag to update other data at each new input)
40
- batch_ptr: tracks batch-id handled by the Triton program
41
- token_chunk_offset_ptr: tracks token group_idx handled by the Triton program
42
- (Triton implementation of causal_conv1d handles parallelism in 3-axes
43
- - feature-axis
44
- - batch-axis
45
- - sequence-axis)
46
- """
47
- nums_dict: Optional[dict] = None
48
- cu_seqlen: Optional[int] = None
49
- batch_ptr: Optional[torch.Tensor] = None
50
- token_chunk_offset_ptr: Optional[torch.Tensor] = None
51
-
52
-
53
-def get_platform_metadata_classes() -> tuple[type[AttentionMetadata], ...]:
54
- """Returns the appropriate metadata classes for the current platform."""
55
- if current_platform.is_rocm():
56
- from vllm.v1.attention.backends.rocm_aiter_fa import (
57
- AiterFlashAttentionMetadata)
58
- from vllm.v1.attention.backends.triton_attn import (
59
- TritonAttentionMetadata)
60
- return (AiterFlashAttentionMetadata, TritonAttentionMetadata,
61
- PlaceholderAttentionMetadata)
62
- if current_platform.is_cuda():
63
- from vllm.v1.attention.backends.flash_attn import (
64
- FlashAttentionMetadata)
65
- from vllm.v1.attention.backends.xformers import (
66
- XFormersAttentionMetadata)
67
- return (FlashAttentionMetadata, XFormersAttentionMetadata,
68
- PlaceholderAttentionMetadata)
69
- raise ValueError(
70
- f"Unsupported platform for Mamba2: {current_platform.device_type}")
71
-
72
-
73
-def prepare_mamba2_metadata(
74
- chunk_size: int,
75
- attn_metadata: AttentionMetadata,
76
-) -> Mamba2Metadata:
77
-
78
- # compute number of prefill and decode requests
79
- # NOTE: in V0 we assume prefills are before decodes
80
- num_prefills = attn_metadata.num_prefills
81
- num_prefill_tokens = attn_metadata.num_prefill_tokens
82
-
83
- seq_idx_p = None
84
- chunk_indices_p, chunk_offsets_p = None, None
85
- # Need flags to indicate if there are initial states
86
- # currently we really only support the FlashAttention backend
87
- has_initial_states_p = None
88
- prep_initial_states = False
89
-
90
- # Compute seq_idx, chunk_indices and chunk_offsets for prefill only
91
- if num_prefills > 0:
92
- attn_metadata_instances = get_platform_metadata_classes()
93
- if (isinstance(attn_metadata, attn_metadata_instances)
94
- and attn_metadata.context_lens_tensor is not None):
95
- # precompute flag to avoid device syncs later in mamba2 layer
96
- # forwards
97
- # prep is only needed for mamba2 ssd prefill processing
98
- has_initial_states_p = (
99
- attn_metadata.context_lens_tensor[:num_prefills] > 0)
100
- prep_initial_states = torch.any(has_initial_states_p).item()
101
- query_start_loc_p = attn_metadata.query_start_loc[:num_prefills + 1]
102
- seq_idx_p = torch.repeat_interleave(torch.arange(
103
- num_prefills, dtype=torch.int32, device=query_start_loc_p.device),
104
- query_start_loc_p.diff(),
105
- output_size=num_prefill_tokens)
106
- seq_idx_p.unsqueeze_(0)
107
-
108
- # We compute metadata for chunked prefill once at the top level model
109
- # forward and reuse them in mamba layers. If not needed, they will be
110
- # ignored inside mamba kernels.
111
- if prep_initial_states:
112
- chunk_indices_p, chunk_offsets_p = \
113
- _query_start_loc_to_chunk_indices_offsets(
114
- query_start_loc_p, chunk_size, num_prefill_tokens)
115
-
116
- return Mamba2Metadata(has_initial_states_p=has_initial_states_p,
117
- prep_initial_states=prep_initial_states,
118
- chunk_size=chunk_size,
119
- seq_idx_p=seq_idx_p,
120
- chunk_indices_p=chunk_indices_p,
121
- chunk_offsets_p=chunk_offsets_p)
122
-
123
-
124
-def update_metadata(x: torch.Tensor, query_start_loc: torch.Tensor,
125
- mamba2_metadata: Union[Mamba2Metadata,
126
- Mamba2AttentionMetadata,
127
- GDNAttentionMetadata]):
128
- """
129
- this is triggered upon handling a new input at the first layer
130
- """
131
- dim, cu_seqlen = x.shape
132
- mamba2_metadata.cu_seqlen = cu_seqlen
133
- seqlens = np.diff(query_start_loc.to('cpu'))
134
- nums_dict = {} # type: ignore
135
- for BLOCK_M in [8]: # cover all BLOCK_M values
136
- nums = -(-seqlens // BLOCK_M)
137
- nums_dict[BLOCK_M] = {}
138
- nums_dict[BLOCK_M]['nums'] = nums
139
- nums_dict[BLOCK_M]['tot'] = nums.sum().item()
140
- mlist = torch.from_numpy(np.repeat(np.arange(len(nums)), nums))
141
- nums_dict[BLOCK_M]['mlist'] = mlist
142
- mlist_len = len(nums_dict[BLOCK_M]['mlist'])
143
- nums_dict[BLOCK_M]['mlist_len'] = mlist_len
144
- MAX_NUM_PROGRAMS = max(1024, mlist_len) * 2
145
- offsetlist = [] # type: ignore
146
- for idx, num in enumerate(nums):
147
- offsetlist.extend(range(num))
148
- offsetlist = torch.tensor(offsetlist, dtype=torch.int32)
149
- nums_dict[BLOCK_M]['offsetlist'] = offsetlist
150
-
151
- if mamba2_metadata.batch_ptr is None:
152
- # Update default value after class definition
153
- #mamba2_metadata.MAX_NUM_PROGRAMS *= 2
154
- mamba2_metadata.batch_ptr = torch.full((MAX_NUM_PROGRAMS, ),
155
- PAD_SLOT_ID,
156
- dtype=torch.int32,
157
- device='cuda')
158
- mamba2_metadata.token_chunk_offset_ptr = torch.full(
159
- (MAX_NUM_PROGRAMS, ),
160
- PAD_SLOT_ID,
161
- dtype=torch.int32,
162
- device='cuda')
163
- else:
164
- if mamba2_metadata.batch_ptr.nelement() < MAX_NUM_PROGRAMS:
165
- mamba2_metadata.batch_ptr.resize_(MAX_NUM_PROGRAMS).fill_(
166
- PAD_SLOT_ID)
167
- mamba2_metadata.token_chunk_offset_ptr.resize_( # type: ignore
168
- MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID)
169
-
170
- mamba2_metadata.batch_ptr[0:mlist_len].copy_(mlist)
171
- mamba2_metadata.token_chunk_offset_ptr[ # type: ignore
172
- 0:mlist_len].copy_(offsetlist)
173
- nums_dict[BLOCK_M]['batch_ptr'] = mamba2_metadata.batch_ptr
174
- nums_dict[BLOCK_M]['token_chunk_offset_ptr'] = (
175
- mamba2_metadata.token_chunk_offset_ptr) # type: ignore
176
- mamba2_metadata.nums_dict = nums_dict
177
- return mamba2_metadata
vllm/model_executor/layers/mamba/mamba_mixer.py
+50
-99
@@ -10,8 +10,6 @@ import torch
10
from torch import nn
11
from torch.nn.parameter import Parameter
12
13
-from vllm import envs
14
-from vllm.attention.backends.abstract import AttentionMetadata
13
from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
14
from vllm.distributed.parallel_state import (
15
get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size)
@@ -28,7 +26,6 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
26
causal_conv1d_fn, causal_conv1d_update)
27
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
28
selective_scan_fn, selective_state_update)
31
-from vllm.model_executor.models.mamba_cache import MambaCacheParams
29
from vllm.model_executor.utils import set_weight_attrs
30
from vllm.platforms import current_platform
31
from vllm.utils import direct_register_custom_op
@@ -149,16 +146,12 @@ class MambaMixer(MambaBase, CustomOp):
146
has_weight=rms_norm_has_weight,
147
) if use_rms_norm else None
148
152
- if envs.VLLM_USE_V1:
153
- compilation_config = get_current_vllm_config().compilation_config
154
- if prefix in compilation_config.static_forward_context:
155
- raise ValueError(f"Duplicate layer name: {prefix}")
156
- compilation_config.static_forward_context[prefix] = self
157
- # The outer list is for v0 PP virtual engine. Though this code path
158
- # only runs for v1, we have to do this to unify with the interface
159
- # of Attention + v0 PP.
160
- # The inner tuple is (conv_state, ssm_state)
161
- self.kv_cache = [(torch.tensor([]), torch.tensor([]))]
149
+ compilation_config = get_current_vllm_config().compilation_config
150
+ if prefix in compilation_config.static_forward_context:
151
+ raise ValueError(f"Duplicate layer name: {prefix}")
152
+ compilation_config.static_forward_context[prefix] = self
153
+ # The inner tuple is (conv_state, ssm_state)
154
+ self.kv_cache = (torch.tensor([]), torch.tensor([]))
155
156
self.model_config = model_config
157
self.cache_config = cache_config
@@ -186,29 +179,18 @@ class MambaMixer(MambaBase, CustomOp):
179
discrete_time_step = self.dt_proj(time_step)[0].transpose(-2, -1)
180
return discrete_time_step, B, C
181
189
- def forward(self,
190
- hidden_states: torch.Tensor,
191
- output: torch.Tensor,
192
- mamba_cache_params: Optional[MambaCacheParams] = None):
193
- if not envs.VLLM_USE_V1:
194
- CustomOp.forward(self, hidden_states, output, mamba_cache_params)
195
- else:
196
- torch.ops.vllm.mamba_mixer(
197
- hidden_states,
198
- output,
199
- self.prefix,
200
- )
201
-
202
- def forward_native(self,
203
- hidden_states: torch.Tensor,
204
- output: torch.Tensor,
205
- mamba_cache_params: Optional[MambaCacheParams] = None):
182
+ def forward(self, hidden_states: torch.Tensor, output: torch.Tensor):
183
+ torch.ops.vllm.mamba_mixer(
184
+ hidden_states,
185
+ output,
186
+ self.prefix,
187
+ )
188
+
189
+ def forward_native(self, hidden_states: torch.Tensor,
190
+ output: torch.Tensor):
191
pass
192
208
- def forward_cuda(self,
209
- hidden_states: torch.Tensor,
210
- output: torch.Tensor,
211
- mamba_cache_params: Optional[MambaCacheParams] = None):
193
+ def forward_cuda(self, hidden_states: torch.Tensor, output: torch.Tensor):
194
"""
195
Run the Mamba-1 SSM pipeline.
196
@@ -234,31 +216,18 @@ class MambaMixer(MambaBase, CustomOp):
216
forward_context: ForwardContext = get_forward_context()
217
attn_metadata = forward_context.attn_metadata
218
237
- if envs.VLLM_USE_V1:
238
- if attn_metadata is not None:
239
- assert isinstance(attn_metadata, dict)
240
- attn_metadata = attn_metadata[self.prefix]
241
- mamba1_metadata = attn_metadata
242
- assert isinstance(mamba1_metadata, Mamba1AttentionMetadata)
243
- query_start_loc = mamba1_metadata.query_start_loc
244
- state_indices_tensor = mamba1_metadata.state_indices_tensor
245
- self_kv_cache = self.kv_cache[forward_context.virtual_engine]
246
- conv_state = self_kv_cache[0].transpose(-1, -2)
247
- ssm_state = self_kv_cache[1]
248
- has_initial_states = mamba1_metadata.has_initial_states
249
- num_padded_decodes = mamba1_metadata.num_padded_decodes
250
- else:
251
- assert isinstance(attn_metadata, AttentionMetadata)
252
- assert mamba_cache_params is not None
253
- conv_state = mamba_cache_params.conv_state
254
- ssm_state = mamba_cache_params.ssm_state
255
- state_indices_tensor = mamba_cache_params.state_indices_tensor
256
- query_start_loc = attn_metadata.query_start_loc
257
- context_lens_tensor = attn_metadata.context_lens_tensor
258
- has_initial_states = None
259
- if context_lens_tensor is not None:
260
- has_initial_states = context_lens_tensor > 0
261
- num_padded_decodes = attn_metadata.num_decode_tokens
219
+ if attn_metadata is not None:
220
+ assert isinstance(attn_metadata, dict)
221
+ attn_metadata = attn_metadata[self.prefix]
222
+ mamba1_metadata = attn_metadata
223
+ assert isinstance(mamba1_metadata, Mamba1AttentionMetadata)
224
+ query_start_loc = mamba1_metadata.query_start_loc
225
+ state_indices_tensor = mamba1_metadata.state_indices_tensor
226
+ self_kv_cache = self.kv_cache[forward_context.virtual_engine]
227
+ conv_state = self_kv_cache[0].transpose(-1, -2)
228
+ ssm_state = self_kv_cache[1]
229
+ has_initial_states = mamba1_metadata.has_initial_states
230
+ num_padded_decodes = mamba1_metadata.num_padded_decodes
231
232
# 1. Gated MLP's linear projection
233
projected_states = self.in_proj(hidden_states)[0].transpose(-2, -1)
@@ -267,7 +236,7 @@ class MambaMixer(MambaBase, CustomOp):
236
conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),
237
self.conv1d.weight.size(2))
238
270
- if envs.VLLM_USE_V1 and attn_metadata is None:
239
+ if attn_metadata is None:
240
# V1 profile run
241
hidden_states_BC = hidden_states_BC.contiguous()
242
return self.out_proj(hidden_states_BC.transpose(-2, -1))[0]
@@ -368,10 +337,7 @@ class MambaMixer(MambaBase, CustomOp):
337
out=scan_outputs_d)
338
scan_outputs_d = scan_outputs_d.transpose(0, 1)
339
371
- if envs.VLLM_USE_V1:
372
- ssm_outputs.insert(0, scan_outputs_d)
373
- else:
374
- ssm_outputs.append(scan_outputs_d)
340
+ ssm_outputs.insert(0, scan_outputs_d)
341
342
scan_outputs_combined = ssm_outputs[0] if len(
343
ssm_outputs) == 1 else torch.cat(ssm_outputs, dim=-1)
@@ -441,40 +407,27 @@ def split_batch_to_prefill_and_decode(
407
num_decodes: int,
408
num_padded_decodes: int,
409
) -> PrefillDecodeSplit:
410
+
411
num_actual_tokens = num_prefill_tokens + num_padded_decodes
412
446
- if envs.VLLM_USE_V1:
447
- # In v1, decode tokens come first, then prefill tokens.
448
- hidden_states_BC_d, hidden_states_BC_p = torch.split(
449
- hidden_states_BC[..., :num_actual_tokens],
450
- [num_padded_decodes, num_prefill_tokens],
451
- dim=-1)
452
- gate_d, gate_p = torch.split(gate[..., :num_actual_tokens],
453
- [num_padded_decodes, num_prefill_tokens],
454
- dim=-1)
455
-
456
- # num_padded_decodes accounts for CUDA graph padding when applicable
457
- state_indices_tensor_d, state_indices_tensor_p = torch.split(
458
- state_indices_tensor[:num_padded_decodes + num_prefills],
459
- [num_padded_decodes, num_prefills],
460
- dim=0)
461
- query_start_loc_p = (query_start_loc[-num_prefills - 1:] -
462
- num_padded_decodes if num_prefills > 0 else None)
463
- has_initial_states_p = has_initial_states[-num_prefills:] if (
464
- has_initial_states is not None and num_prefills > 0) else None
465
- else:
466
- # In v0, prefill tokens come first, then decode tokens.
467
- hidden_states_BC_p, hidden_states_BC_d = torch.split(
468
- hidden_states_BC, [num_prefill_tokens, num_decode_tokens], dim=-1)
469
- gate_p, gate_d = torch.split(gate,
470
- [num_prefill_tokens, num_decode_tokens],
471
- dim=-1)
472
- state_indices_tensor_p, state_indices_tensor_d = torch.split(
473
- state_indices_tensor, [num_prefills, num_decodes], dim=0)
474
- query_start_loc_p = (query_start_loc[:num_prefills +
475
- 1] if num_prefills > 0 else None)
476
- has_initial_states_p = has_initial_states[:num_prefills] if (
477
- has_initial_states is not None and num_prefills > 0) else None
413
+ # In v1, decode tokens come first, then prefill tokens.
414
+ hidden_states_BC_d, hidden_states_BC_p = torch.split(
415
+ hidden_states_BC[..., :num_actual_tokens],
416
+ [num_padded_decodes, num_prefill_tokens],
417
+ dim=-1)
418
+ gate_d, gate_p = torch.split(gate[..., :num_actual_tokens],
419
+ [num_padded_decodes, num_prefill_tokens],
420
+ dim=-1)
421
+
422
+ # num_padded_decodes accounts for CUDA graph padding when applicable
423
+ state_indices_tensor_d, state_indices_tensor_p = torch.split(
424
+ state_indices_tensor[:num_padded_decodes + num_prefills],
425
+ [num_padded_decodes, num_prefills],
426
+ dim=0)
427
+ query_start_loc_p = (query_start_loc[-num_prefills - 1:] -
428
+ num_padded_decodes if num_prefills > 0 else None)
429
+ has_initial_states_p = has_initial_states[-num_prefills:] if (
430
+ has_initial_states is not None and num_prefills > 0) else None
431
432
return PrefillDecodeSplit(
433
hidden_states_BC_p=hidden_states_BC_p,
@@ -495,9 +448,7 @@ def mamba_mixer(
448
) -> None:
449
forward_context: ForwardContext = get_forward_context()
450
self = forward_context.no_compile_layers[layer_name]
498
- self.forward_cuda(hidden_states=hidden_states,
499
- output=output,
500
- mamba_cache_params=None)
451
+ self.forward_cuda(hidden_states=hidden_states, output=output)
452
453
454
def mamba_mixer_fake(
vllm/model_executor/layers/mamba/mamba_mixer2.py
+56
-116
@@ -9,7 +9,6 @@ if TYPE_CHECKING:
9
import torch
10
from torch import nn
11
12
-from vllm import envs
12
from vllm.attention.backends.abstract import AttentionMetadata
13
from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
14
from vllm.distributed import (divide, get_tensor_model_parallel_rank,
@@ -22,8 +21,6 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
21
MergedColumnParallelLinear,
22
RowParallelLinear)
23
from vllm.model_executor.layers.mamba.abstract import MambaBase
25
-from vllm.model_executor.layers.mamba.mamba2_metadata import (Mamba2Metadata,
26
- update_metadata)
24
from vllm.model_executor.layers.mamba.mamba_utils import (
25
MambaStateDtypeCalculator, MambaStateShapeCalculator)
26
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
@@ -36,7 +33,6 @@ from vllm.model_executor.layers.mamba.ops.ssd_combined import (
33
from vllm.model_executor.layers.quantization import QuantizationConfig
34
from vllm.model_executor.model_loader.weight_utils import (
35
LoaderFunction, composed_weight_loader, sharded_weight_loader)
39
-from vllm.model_executor.models.mamba_cache import MambaCacheParams
36
from vllm.model_executor.utils import set_weight_attrs
37
from vllm.platforms import current_platform
38
from vllm.utils import direct_register_custom_op
@@ -449,16 +445,12 @@ class MambaMixer2(MambaBase, CustomOp):
445
self.use_rms_norm,
446
eps=rms_norm_eps)
447
452
- if envs.VLLM_USE_V1:
453
- compilation_config = get_current_vllm_config().compilation_config
454
- if prefix in compilation_config.static_forward_context:
455
- raise ValueError(f"Duplicate layer name: {prefix}")
456
- compilation_config.static_forward_context[prefix] = self
457
- # The outer list is for v0 PP virtual engine. Though this code path
458
- # only runs for v1, we have to do this to unify with the interface
459
- # of Attention + v0 PP.
460
- # The inner tuple is (conv_state, ssm_state)
461
- self.kv_cache = [(torch.tensor([]), torch.tensor([]))]
448
+ compilation_config = get_current_vllm_config().compilation_config
449
+ if prefix in compilation_config.static_forward_context:
450
+ raise ValueError(f"Duplicate layer name: {prefix}")
451
+ compilation_config.static_forward_context[prefix] = self
452
+ # The tuple is (conv_state, ssm_state)
453
+ self.kv_cache = (torch.tensor([]), torch.tensor([]))
454
455
self.model_config = model_config
456
self.cache_config = cache_config
@@ -468,8 +460,6 @@ class MambaMixer2(MambaBase, CustomOp):
460
self,
461
hidden_states: torch.Tensor,
462
output: torch.Tensor,
471
- mamba_cache_params: MambaCacheParams,
472
- mamba2_metadata: Mamba2Metadata,
463
mup_vector: Optional[torch.Tensor] = None,
464
):
465
pass
@@ -478,59 +468,43 @@ class MambaMixer2(MambaBase, CustomOp):
468
self,
469
hidden_states: torch.Tensor,
470
output: torch.Tensor,
481
- mamba_cache_params: MambaCacheParams,
482
- mamba2_metadata: Mamba2Metadata,
471
mup_vector: Optional[torch.Tensor] = None,
472
):
485
- if not envs.VLLM_USE_V1:
486
- CustomOp.forward(self, hidden_states, output, mamba_cache_params,
487
- mamba2_metadata, mup_vector)
488
- else:
489
- torch.ops.vllm.mamba_mixer2(
490
- hidden_states,
491
- output,
492
- self.prefix,
493
- mup_vector,
494
- )
473
+ torch.ops.vllm.mamba_mixer2(
474
+ hidden_states,
475
+ output,
476
+ self.prefix,
477
+ mup_vector,
478
+ )
479
480
def forward_cuda(
481
self,
482
hidden_states: torch.Tensor,
483
output: torch.Tensor,
500
- mamba_cache_params: MambaCacheParams,
501
- mamba2_metadata: Mamba2Metadata,
484
mup_vector: Optional[torch.Tensor] = None,
485
):
486
forward_context = get_forward_context()
505
- # mamba2_metadata contains metadata necessary for the mamba2 triton
487
+ # attn_metadata contains metadata necessary for the mamba2 triton
488
# kernels to operate in continuous batching and in chunked prefill
489
# modes; they are computed at top-level model forward since they
490
# stay the same and reused for all mamba layers in the same iteration
491
attn_metadata: AttentionMetadata = forward_context.attn_metadata
510
- if envs.VLLM_USE_V1:
511
- if attn_metadata is not None:
512
- assert isinstance(attn_metadata, dict)
513
- attn_metadata = attn_metadata[self.prefix]
514
- mamba2_metadata = attn_metadata
515
- assert isinstance(attn_metadata, Mamba2AttentionMetadata)
516
- self_kv_cache = self.kv_cache[forward_context.virtual_engine]
517
- # conv_state = (..., dim, width-1) yet contiguous along 'dim'
518
- conv_state = self_kv_cache[0].transpose(-1, -2)
519
- ssm_state = self_kv_cache[1]
520
- state_indices_tensor = attn_metadata.state_indices_tensor
521
- else:
522
- conv_state = mamba_cache_params.conv_state
523
- ssm_state = mamba_cache_params.ssm_state
524
- state_indices_tensor = mamba_cache_params.state_indices_tensor
525
-
526
- # Common members between V1 metadata and V0 metadata
527
- if mamba2_metadata is not None:
528
- has_initial_states_p = mamba2_metadata.has_initial_states_p
529
- prep_initial_states = mamba2_metadata.prep_initial_states
530
- chunk_size = mamba2_metadata.chunk_size
531
- seq_idx_p = mamba2_metadata.seq_idx_p
532
- chunk_indices_p = mamba2_metadata.chunk_indices_p
533
- chunk_offsets_p = mamba2_metadata.chunk_offsets_p
492
+
493
+ if attn_metadata is not None:
494
+ assert isinstance(attn_metadata, dict)
495
+ attn_metadata = attn_metadata[self.prefix]
496
+ assert isinstance(attn_metadata, Mamba2AttentionMetadata)
497
+ self_kv_cache = self.kv_cache[forward_context.virtual_engine]
498
+ # conv_state = (..., dim, width-1) yet contiguous along 'dim'
499
+ conv_state = self_kv_cache[0].transpose(-1, -2)
500
+ ssm_state = self_kv_cache[1]
501
+ state_indices_tensor = attn_metadata.state_indices_tensor
502
+ has_initial_states_p = attn_metadata.has_initial_states_p
503
+ prep_initial_states = attn_metadata.prep_initial_states
504
+ chunk_size = attn_metadata.chunk_size
505
+ seq_idx_p = attn_metadata.seq_idx_p
506
+ chunk_indices_p = attn_metadata.chunk_indices_p
507
+ chunk_offsets_p = attn_metadata.chunk_offsets_p
508
509
# 1. Gated MLP's linear projection
510
projected_states, _ = self.in_proj(hidden_states)
@@ -562,8 +536,8 @@ class MambaMixer2(MambaBase, CustomOp):
536
dim=-1,
537
)
538
565
- if envs.VLLM_USE_V1 and attn_metadata is None:
566
- # V1 profile run
539
+ if attn_metadata is None:
540
+ # profile run
541
hidden_states_B_C = (hidden_states_B_C.transpose(
542
0, 1).clone().transpose(0, 1)).contiguous()
543
hidden_states, _B, _C = split_hidden_states_B_C_fn(
@@ -579,49 +553,27 @@ class MambaMixer2(MambaBase, CustomOp):
553
has_decode = num_decodes > 0
554
num_actual_tokens = num_prefill_tokens + num_decodes
555
582
- # NOTE: V0 put prefill before decode, v1 puts decode before prefill
556
# Separate prefill and decode by splitting varlen input
557
# Split along token dimension
585
- if envs.VLLM_USE_V1:
586
- hidden_states_B_C_d, hidden_states_B_C_p = torch.split(
587
- hidden_states_B_C[:num_actual_tokens],
588
- [num_decodes, num_prefill_tokens],
589
- dim=0,
590
- )
591
- dt_d, dt_p = torch.split(
592
- dt[:num_actual_tokens],
593
- [num_decodes, num_prefill_tokens],
594
- dim=0,
595
- )
596
- # Split along batch dimension
597
- state_indices_tensor_d, state_indices_tensor_p = torch.split(
598
- state_indices_tensor[:num_actual_tokens],
599
- [num_decodes, num_prefills],
600
- dim=0,
601
- )
602
- query_start_loc_p = (
603
- attn_metadata.query_start_loc[-num_prefills - 1:] -
604
- num_decodes if has_prefill else None)
605
- else:
606
- hidden_states_B_C_p, hidden_states_B_C_d = torch.split(
607
- hidden_states_B_C,
608
- [num_prefill_tokens, num_decodes],
609
- dim=0,
610
- )
611
- dt_p, dt_d = torch.split(
612
- dt,
613
- [num_prefill_tokens, num_decodes],
614
- dim=0,
615
- )
616
- # Split along batch dimension
617
- state_indices_tensor_p, state_indices_tensor_d = torch.split(
618
- state_indices_tensor,
619
- [num_prefills, num_decodes],
620
- dim=0,
621
- )
622
- query_start_loc_p = (attn_metadata.query_start_loc[:num_prefills +
623
- 1]
624
- if has_prefill else None)
558
+ hidden_states_B_C_d, hidden_states_B_C_p = torch.split(
559
+ hidden_states_B_C[:num_actual_tokens],
560
+ [num_decodes, num_prefill_tokens],
561
+ dim=0,
562
+ )
563
+ dt_d, dt_p = torch.split(
564
+ dt[:num_actual_tokens],
565
+ [num_decodes, num_prefill_tokens],
566
+ dim=0,
567
+ )
568
+ # Split along batch dimension
569
+ state_indices_tensor_d, state_indices_tensor_p = torch.split(
570
+ state_indices_tensor[:num_actual_tokens],
571
+ [num_decodes, num_prefills],
572
+ dim=0,
573
+ )
574
+ query_start_loc_p = (
575
+ attn_metadata.query_start_loc[-num_prefills - 1:] -
576
+ num_decodes if has_prefill else None)
577
578
# Preallocate output tensor to avoid memcpy cost for merging prefill
579
# and decode outputs
@@ -633,18 +585,11 @@ class MambaMixer2(MambaBase, CustomOp):
585
dtype=hidden_states.dtype,
586
device=hidden_states.device,
587
)
636
- if envs.VLLM_USE_V1:
637
- preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(
638
- preallocated_ssm_out,
639
- [num_decodes, num_prefill_tokens],
640
- dim=0,
641
- )
642
- else:
643
- preallocated_ssm_out_p, preallocated_ssm_out_d = torch.split(
644
- preallocated_ssm_out,
645
- [num_prefill_tokens, num_decodes],
646
- dim=0,
647
- )
588
+ preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(
589
+ preallocated_ssm_out,
590
+ [num_decodes, num_prefill_tokens],
591
+ dim=0,
592
+ )
593
594
# Process prefill requests
595
if has_prefill:
@@ -653,9 +598,6 @@ class MambaMixer2(MambaBase, CustomOp):
598
# pointed to by "state_indices_tensor"
599
x = hidden_states_B_C_p.transpose(
600
0, 1) # this is the form that causal-conv see
656
- if mamba2_metadata.cu_seqlen is None:
657
- mamba2_metadata = update_metadata(x, query_start_loc_p,
658
- mamba2_metadata)
601
hidden_states_B_C_p = causal_conv1d_fn(
602
x,
603
conv_weights,
@@ -664,7 +606,7 @@ class MambaMixer2(MambaBase, CustomOp):
606
conv_states=conv_state,
607
has_initial_state=has_initial_states_p,
608
cache_indices=state_indices_tensor_p,
667
- metadata=mamba2_metadata,
609
+ metadata=attn_metadata,
610
query_start_loc=query_start_loc_p).transpose(
611
0, 1)[:num_prefill_tokens]
612
@@ -806,8 +748,6 @@ def mamba_mixer2(
748
self = forward_context.no_compile_layers[layer_name]
749
self.forward_cuda(hidden_states=hidden_states,
750
output=output,
809
- mamba_cache_params=None,
810
- mamba2_metadata=None,
751
mup_vector=mup_vector)
752
753
vllm/model_executor/layers/mamba/mamba_utils.py
+2
-18
@@ -100,7 +100,6 @@ class MambaStateShapeCalculator:
100
intermediate_size: int,
101
state_size: int,
102
conv_kernel: int,
103
- use_v1: bool = True,
103
) -> tuple[tuple[int, int], tuple[int, int]]:
104
conv_state_shape = (divide(intermediate_size,
105
tp_world_size), conv_kernel - 1)
@@ -108,11 +107,7 @@ class MambaStateShapeCalculator:
107
temporal_state_shape = (divide(intermediate_size,
108
tp_world_size), state_size)
109
111
- # In V0, the conv_state shape was swapped during allocation in
112
- # MambaCacheManager, but in V1 it needs to be determined here at the
113
- # calculation level
114
- if use_v1:
115
- conv_state_shape = conv_state_shape[1], conv_state_shape[0]
110
+ conv_state_shape = conv_state_shape[1], conv_state_shape[0]
111
112
return conv_state_shape, temporal_state_shape
113
@@ -126,7 +121,6 @@ class MambaStateShapeCalculator:
121
head_dim: int,
122
state_size: int,
123
conv_kernel: int,
129
- use_v1: bool = True,
124
) -> tuple[tuple[int, int], tuple[int, int, int]]:
125
# if n_groups is not divisible by world_size, need to extend the shards
126
# to ensure all groups needed by a head is sharded along with it
@@ -137,8 +131,6 @@ class MambaStateShapeCalculator:
131
132
# contiguous along 'dim' axis
133
conv_state_shape = (conv_kernel - 1, divide(conv_dim, tp_world_size))
140
- if not use_v1:
141
- conv_state_shape = conv_state_shape[1], conv_state_shape[0]
134
135
# These are not TP-ed as they depend on A, dt_bias, D
136
# - they are typically small
@@ -153,12 +145,9 @@ class MambaStateShapeCalculator:
145
tp_world_size: int,
146
intermediate_size: int,
147
conv_kernel: int,
156
- use_v1: bool = True,
148
) -> tuple[tuple[int, int]]:
149
conv_dim = divide(intermediate_size, tp_world_size)
150
conv_state_shape = (conv_kernel - 1, conv_dim)
160
- if not use_v1:
161
- conv_state_shape = conv_state_shape[1], conv_state_shape[0]
151
return (conv_state_shape, )
152
153
@classmethod
@@ -183,7 +172,6 @@ class MambaStateShapeCalculator:
172
head_v_dim: int,
173
conv_kernel_size: int,
174
num_spec: int = 0,
186
- use_v1: bool = True,
175
):
176
conv_dim = (head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads)
177
conv_state_shape = (
@@ -191,11 +179,7 @@ class MambaStateShapeCalculator:
179
conv_kernel_size - 1 + num_spec,
180
)
181
194
- # In V0, the conv_state shape was swapped during allocation in
195
- # MambaCacheManager, but in V1 it needs to be determined here at the
196
- # calculation level
197
- if use_v1:
198
- conv_state_shape = conv_state_shape[1], conv_state_shape[0]
182
+ conv_state_shape = conv_state_shape[1], conv_state_shape[0]
183
184
temporal_state_shape = (divide(num_v_heads,
185
tp_world_size), head_k_dim, head_v_dim)
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
-3
@@ -420,9 +420,7 @@ def causal_conv1d_fn(
420
x = x.to(conv_states.dtype)
421
out = torch.empty_like(x)
422
if metadata is not None:
423
- cu_seqlen = metadata.cu_seqlen
423
nums_dict = metadata.nums_dict
425
- #x = metadata.x
424
args = nums_dict
425
batch_ptr = metadata.batch_ptr
426
token_chunk_offset_ptr = metadata.token_chunk_offset_ptr
@@ -926,7 +924,6 @@ def causal_conv1d_update(
924
query_start_loc: Optional[torch.Tensor] = None,
925
max_query_len: int = -1,
926
pad_slot_id: int = PAD_SLOT_ID,
929
- metadata=None,
927
validate_data=False,
928
):
929
"""
vllm/model_executor/layers/mamba/short_conv.py
+3
-18
@@ -8,7 +8,6 @@ if TYPE_CHECKING:
8
9
import torch
10
11
-from vllm import envs
11
from vllm.attention.backends.abstract import AttentionMetadata
12
from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
13
from vllm.distributed import get_tensor_model_parallel_world_size
@@ -18,7 +17,6 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
17
MergedColumnParallelLinear,
18
RowParallelLinear)
19
from vllm.model_executor.layers.mamba.abstract import MambaBase
21
-from vllm.model_executor.layers.mamba.mamba2_metadata import update_metadata
20
from vllm.model_executor.layers.mamba.mamba_utils import (
21
MambaStateDtypeCalculator, MambaStateShapeCalculator)
22
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
@@ -71,15 +69,11 @@ class ShortConv(MambaBase, CustomOp):
69
prefix=f"{prefix}.out_proj",
70
)
71
74
- assert envs.VLLM_USE_V1, ("ShortConv layers are only supported in V1")
72
compilation_config = get_current_vllm_config().compilation_config
73
if prefix in compilation_config.static_forward_context:
74
raise ValueError(f"Duplicate layer name: {prefix}")
75
compilation_config.static_forward_context[prefix] = self
79
- # The outer list is for v0 PP virtual engine. Though this code path
80
- # only runs for v1, we have to do this to unify with the interface
81
- # of Attention + v0 PP.
82
- self.kv_cache = [(torch.tensor([]), )]
76
+ self.kv_cache = (torch.tensor([]), )
77
78
self.model_config = model_config
79
self.cache_config = cache_config
@@ -89,7 +83,6 @@ class ShortConv(MambaBase, CustomOp):
83
self,
84
hidden_states: torch.Tensor,
85
output: torch.Tensor,
92
- conv_metadata: ShortConvAttentionMetadata,
86
):
87
return
88
@@ -97,7 +90,6 @@ class ShortConv(MambaBase, CustomOp):
90
self,
91
hidden_states: torch.Tensor,
92
output: torch.Tensor,
100
- conv_metadata: ShortConvAttentionMetadata,
93
):
94
torch.ops.vllm.short_conv(
95
hidden_states,
@@ -109,7 +101,6 @@ class ShortConv(MambaBase, CustomOp):
101
self,
102
hidden_states: torch.Tensor,
103
output: torch.Tensor,
112
- conv_metadata: ShortConvAttentionMetadata,
104
):
105
forward_context = get_forward_context()
106
# ShortConvAttentionMetadata contains metadata necessary for the
@@ -121,7 +112,6 @@ class ShortConv(MambaBase, CustomOp):
112
if attn_metadata is not None:
113
assert isinstance(attn_metadata, dict)
114
attn_metadata = attn_metadata[self.prefix]
124
- conv_metadata = attn_metadata
115
assert isinstance(attn_metadata, ShortConvAttentionMetadata)
116
self_kv_cache = self.kv_cache[forward_context.virtual_engine]
117
conv_state = self_kv_cache[0].transpose(-1, -2)
@@ -181,9 +171,6 @@ class ShortConv(MambaBase, CustomOp):
171
172
if has_prefill:
173
Bx_p = (B_p * x_p).transpose(0, 1)
184
- if conv_metadata.cu_seqlen is None:
185
- conv_metadata = update_metadata(Bx_p, query_start_loc_p,
186
- conv_metadata)
174
Bx = causal_conv1d_fn(Bx_p,
175
conv_weights,
176
self.conv.bias,
@@ -191,7 +178,7 @@ class ShortConv(MambaBase, CustomOp):
178
conv_states=conv_state,
179
has_initial_state=has_initial_states_p,
180
cache_indices=state_indices_tensor_p,
194
- metadata=conv_metadata,
181
+ metadata=attn_metadata,
182
query_start_loc=query_start_loc_p).transpose(
183
0, 1)[:num_prefill_tokens]
184
@@ -248,9 +235,7 @@ def short_conv(
235
) -> None:
236
forward_context: ForwardContext = get_forward_context()
237
self = forward_context.no_compile_layers[layer_name]
251
- self.forward_cuda(hidden_states=hidden_states,
252
- output=output,
253
- conv_metadata=None)
238
+ self.forward_cuda(hidden_states=hidden_states, output=output)
239
240
241
def short_conv_fake(
vllm/model_executor/models/bamba.py
+3
-69
@@ -9,21 +9,17 @@ import torch
9
from torch import nn
10
from transformers import BambaConfig
11
12
-from vllm import envs
12
from vllm.attention.layer import Attention
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import CacheConfig, ModelConfig, VllmConfig
15
from vllm.distributed import get_tensor_model_parallel_world_size
16
from vllm.distributed.parallel_state import get_pp_group
18
-from vllm.forward_context import get_forward_context
17
from vllm.model_executor.layers.activation import SiluAndMul
18
from vllm.model_executor.layers.layernorm import RMSNorm
19
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
20
QKVParallelLinear,
21
RowParallelLinear)
22
from vllm.model_executor.layers.logits_processor import LogitsProcessor
25
-from vllm.model_executor.layers.mamba.mamba2_metadata import (
26
- Mamba2Metadata, prepare_mamba2_metadata)
23
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
24
from vllm.model_executor.layers.mamba.mamba_utils import (
25
MambaStateDtypeCalculator, MambaStateShapeCalculator)
@@ -32,10 +28,7 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
28
from vllm.model_executor.layers.vocab_parallel_embedding import (
29
DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
30
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
35
-from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
36
- MambaCacheParams)
31
from vllm.sequence import IntermediateTensors
38
-from vllm.utils import LayerBlockType
32
33
from .interfaces import (HasInnerState, IsHybrid, SupportsLoRA, SupportsPP,
34
SupportsQuant)
@@ -115,8 +108,6 @@ class BambaMixerDecoderLayer(nn.Module):
108
self,
109
hidden_states: torch.Tensor,
110
residual: Optional[torch.Tensor],
118
- mamba_cache_params: MambaCacheParams,
119
- mamba2_metadata: Mamba2Metadata,
111
**kwargs,
112
):
113
if residual is None:
@@ -127,7 +118,7 @@ class BambaMixerDecoderLayer(nn.Module):
118
hidden_states, residual)
119
120
output = torch.empty_like(hidden_states)
130
- self.mamba(hidden_states, output, mamba_cache_params, mamba2_metadata)
121
+ self.mamba(hidden_states, output)
122
# Fully Connected
123
hidden_states, residual = self.pre_ff_layernorm(output, residual)
124
hidden_states = self.feed_forward(hidden_states)
@@ -315,22 +306,10 @@ class BambaModel(nn.Module):
306
self,
307
input_ids: torch.Tensor,
308
positions: torch.Tensor,
318
- mamba_cache_params: MambaCacheParams,
309
intermediate_tensors: Optional[IntermediateTensors] = None,
310
inputs_embeds: Optional[torch.Tensor] = None,
311
) -> torch.Tensor:
312
323
- attn_metadata = get_forward_context().attn_metadata
324
-
325
- if not envs.VLLM_USE_V1:
326
- mamba2_metadata = prepare_mamba2_metadata(
327
- chunk_size=self.config.mamba_chunk_size,
328
- attn_metadata=attn_metadata,
329
- )
330
- else:
331
- # v1 get mamba2_metadata from forward_context
332
- mamba2_metadata = None
333
-
313
if get_pp_group().is_first_rank:
314
if inputs_embeds is not None:
315
hidden_states = inputs_embeds
@@ -343,23 +322,11 @@ class BambaModel(nn.Module):
322
residual = intermediate_tensors["residual"]
323
324
residual = None
346
- num_attn = 0
325
for i, layer in enumerate(self.layers):
348
- if isinstance(layer, BambaAttentionDecoderLayer):
349
- num_attn += 1
350
-
351
- layer_mamba_cache_params = None
352
- if isinstance(layer,
353
- BambaMixerDecoderLayer) and mamba_cache_params:
354
- layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
355
- i - num_attn)
356
-
326
hidden_states, residual = layer(
327
positions=positions,
328
hidden_states=hidden_states,
329
residual=residual,
361
- mamba_cache_params=layer_mamba_cache_params,
362
- mamba2_metadata=mamba2_metadata,
330
)
331
332
if not get_pp_group().is_last_rank:
@@ -457,13 +424,11 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
424
def get_mamba_state_shape_from_config(
425
cls,
426
vllm_config: "VllmConfig",
460
- use_v1: bool = True,
427
) -> tuple[tuple[int, int], tuple[int, int, int]]:
428
"""Calculate shapes for Mamba's convolutional and state caches.
429
430
Args:
431
vllm_config: vLLM config
466
- use_v1: Get shapes for V1 (or V0)
432
433
Returns:
434
Tuple containing:
@@ -482,7 +447,6 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
447
head_dim=hf_config.mamba_d_head,
448
state_size=hf_config.mamba_d_state,
449
conv_kernel=hf_config.mamba_d_conv,
485
- use_v1=use_v1,
450
)
451
452
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -515,8 +479,6 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
479
if not lora_config else lora_config.lora_vocab_padding_size,
480
prefix=maybe_prefix(prefix, "lm_head"),
481
)
518
- # Used to track and store by the Mamba cache between steps.
519
- self.mamba_cache: Optional[MambaCacheManager] = None
482
483
self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
484
config.vocab_size)
@@ -534,39 +496,11 @@ class BambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
496
inputs_embeds: Optional[torch.Tensor] = None,
497
**kwargs):
498
537
- mamba_cache_params = None
538
- if not envs.VLLM_USE_V1:
539
- if self.mamba_cache is None:
540
- num_mamba_layers = \
541
- self.model_config.get_num_layers_by_block_type(
542
- self.vllm_config.parallel_config,
543
- LayerBlockType.mamba
544
- )
545
- mamba_state_shape = \
546
- self.get_mamba_state_shape_from_config(
547
- self.vllm_config, use_v1=False)
548
- mamba_state_dtype = \
549
- self.get_mamba_state_dtype_from_config(
550
- self.vllm_config)
551
- self.mamba_cache = MambaCacheManager(self.vllm_config,
552
- num_mamba_layers,
553
- *mamba_state_shape,
554
- *mamba_state_dtype)
555
-
556
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
557
-
558
- hidden_states = self.model(input_ids, positions, mamba_cache_params,
559
- intermediate_tensors, inputs_embeds)
499
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
500
+ inputs_embeds)
501
502
return hidden_states
503
563
- def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
564
- return self.mamba_cache.copy_inputs_before_cuda_graphs(
565
- input_buffers, **kwargs)
566
-
567
- def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
568
- return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
569
-
504
def compute_logits(
505
self,
506
hidden_states: torch.Tensor,
vllm/model_executor/models/constant_size_cache.py
deleted
-137
@@ -1,137 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from abc import ABC, abstractmethod
4
-from typing import Any
5
-
6
-import torch
7
-
8
-from vllm.attention.backends.utils import PAD_SLOT_ID
9
-
10
-
11
-class ConstantSizeCache(ABC):
12
- """
13
- Abstract base class for managing constant size caches
14
- like Mamba and Minimax.
15
- """
16
-
17
- def __init__(self, max_batch_size: int):
18
- # Maps between the request id and a dict that maps between the seq_id
19
- # and its index inside the cache
20
- self.cache_indices_mapping: dict[str, dict[int, int]] = {}
21
- self.free_cache_indices = list(range(max_batch_size))
22
-
23
- @property
24
- @abstractmethod
25
- def cache(self) -> Any:
26
- """Return the underlying cache tensor(s)"""
27
- pass
28
-
29
- @abstractmethod
30
- def _copy_cache(self, from_index: int, to_index: int):
31
- """Copy cache data from one index to another"""
32
- pass
33
-
34
- def current_run_tensors(self, **kwargs) -> tuple:
35
- """
36
- Return the tensors for the current run's conv and ssm state.
37
- """
38
- if "seqlen_agnostic_capture_inputs" not in kwargs:
39
- # We get here only on Prefill/Eager mode runs
40
- request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"]
41
- finished_requests_ids = kwargs["finished_requests_ids"]
42
-
43
- self._release_finished_requests(finished_requests_ids)
44
- state_indices = self._prepare_current_run_cache(
45
- request_ids_to_seq_ids, finished_requests_ids)
46
-
47
- state_indices_tensor = torch.as_tensor(state_indices,
48
- dtype=torch.int32,
49
- device="cuda")
50
- cache_tensors = self.cache
51
- else:
52
- # CUDA graph capturing runs
53
- cache_tensors, state_indices_tensor = kwargs[
54
- "seqlen_agnostic_capture_inputs"]
55
-
56
- return (cache_tensors, state_indices_tensor)
57
-
58
- def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
59
- """
60
- Copy the relevant state_indices into the CUDA graph input buffer
61
- """
62
- assert all(
63
- key in kwargs
64
- for key in ["request_ids_to_seq_ids", "finished_requests_ids"])
65
- finished_requests_ids = kwargs["finished_requests_ids"]
66
- request_ids_to_seq_ids = kwargs["request_ids_to_seq_ids"]
67
- assert "seqlen_agnostic_capture_inputs" in input_buffers
68
- _, input_state_indices_buffer = input_buffers[
69
- "seqlen_agnostic_capture_inputs"]
70
-
71
- self._release_finished_requests(finished_requests_ids)
72
- state_indices = self._prepare_current_run_cache(
73
- request_ids_to_seq_ids, finished_requests_ids)
74
- cuda_graph_pad_len = input_state_indices_buffer.shape[0] - len(
75
- state_indices)
76
- state_indices.extend([PAD_SLOT_ID] * cuda_graph_pad_len)
77
-
78
- input_state_indices_buffer.copy_(
79
- torch.as_tensor(state_indices, dtype=torch.int32, device="cuda"))
80
-
81
- def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
82
- """
83
- Provide the CUDA graph capture runs with a buffer in adjusted size.
84
- The buffer is used to maintain the Cache during the CUDA graph replay
85
- runs.
86
- """
87
- state_indices_tensor = torch.as_tensor([PAD_SLOT_ID] * batch_size,
88
- dtype=torch.int32,
89
- device="cuda")
90
- return (self.cache, state_indices_tensor)
91
-
92
- def _assign_seq_id_to_cache_index(self, cur_rid: str, seq_id: int,
93
- finished_requests_ids) -> int:
94
- """
95
- Assign (req_id,seq_id) pair to a `destination_index` index, if
96
- already occupied, move the occupying index to a free index.
97
- """
98
- if cur_rid in finished_requests_ids:
99
- # set as pad, do not allocate destination index
100
- return PAD_SLOT_ID
101
- elif cur_rid not in self.cache_indices_mapping:
102
- destination_index = self.free_cache_indices.pop()
103
- self.cache_indices_mapping[cur_rid] = {seq_id: destination_index}
104
- return destination_index
105
- elif seq_id not in (seq_ids2indices :=
106
- self.cache_indices_mapping[cur_rid]):
107
- # parallel sampling , where n > 1, assume prefill have
108
- # already happened, so we copy the
109
- # existing cache into the siblings seq_ids caches
110
- index_exists = next(iter(seq_ids2indices.values()))
111
- # case of decoding n>1, copy prefill cache to decoding indices
112
- destination_index = self.free_cache_indices.pop()
113
- self._copy_cache(from_index=index_exists,
114
- to_index=destination_index)
115
- self.cache_indices_mapping[cur_rid][seq_id] = destination_index
116
- return destination_index
117
- else:
118
- return self.cache_indices_mapping[cur_rid][seq_id]
119
-
120
- def _prepare_current_run_cache(
121
- self, request_ids_to_seq_ids: dict[str, list[int]],
122
- finished_requests_ids: list[str]) -> list[int]:
123
- return [
124
- self._assign_seq_id_to_cache_index(req_id, seq_id,
125
- finished_requests_ids)
126
- for req_id, seq_ids in request_ids_to_seq_ids.items()
127
- for seq_id in seq_ids
128
- ]
129
-
130
- def _release_finished_requests(self,
131
- finished_seq_groups_req_ids: list[str]):
132
- for req_id in finished_seq_groups_req_ids:
133
- if req_id in self.cache_indices_mapping:
134
- for seq_id in self.cache_indices_mapping[req_id]:
135
- self.free_cache_indices.append(
136
- self.cache_indices_mapping[req_id][seq_id])
137
- self.cache_indices_mapping.pop(req_id)
vllm/model_executor/models/falcon_h1.py
+1
-64
@@ -8,21 +8,17 @@ import torch
8
from torch import nn
9
from transformers import FalconH1Config
10
11
-from vllm import envs
11
from vllm.attention.layer import Attention
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, ModelConfig, VllmConfig
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.distributed.parallel_state import get_pp_group
17
-from vllm.forward_context import get_forward_context
16
from vllm.model_executor.layers.activation import SiluAndMul
17
from vllm.model_executor.layers.layernorm import RMSNorm
18
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
19
QKVParallelLinear,
20
RowParallelLinear)
21
from vllm.model_executor.layers.logits_processor import LogitsProcessor
24
-from vllm.model_executor.layers.mamba.mamba2_metadata import (
25
- Mamba2Metadata, prepare_mamba2_metadata)
22
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
23
from vllm.model_executor.layers.mamba.mamba_utils import (
24
MambaStateDtypeCalculator, MambaStateShapeCalculator)
@@ -31,8 +27,6 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
27
from vllm.model_executor.layers.vocab_parallel_embedding import (
28
DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
29
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
34
-from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
35
- MambaCacheParams)
30
from vllm.sequence import IntermediateTensors
31
32
from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP
@@ -179,16 +173,12 @@ class FalconH1SSMDecoderLayer(nn.Module):
173
self,
174
hidden_states: torch.Tensor,
175
residual: Optional[torch.Tensor],
182
- mamba_cache_params: MambaCacheParams,
183
- mamba2_metadata: Mamba2Metadata,
176
**kwargs,
177
):
178
output = torch.empty_like(hidden_states)
179
self.mamba(
180
hidden_states,
181
output,
190
- mamba_cache_params,
191
- mamba2_metadata=mamba2_metadata,
182
mup_vector=self.mup_vector,
183
)
184
return output, residual
@@ -364,8 +354,6 @@ class FalconH1ParallelHybrid(nn.Module):
354
self,
355
positions: torch.Tensor,
356
hidden_states: torch.Tensor,
367
- mamba_cache_params: MambaCacheParams,
368
- mamba2_metadata: Mamba2Metadata,
357
**kwargs,
358
):
359
residual = hidden_states
@@ -382,12 +370,10 @@ class FalconH1ParallelHybrid(nn.Module):
370
371
# Process input through the SSM branch.
372
# FalconH1SSMDecoderLayer expects hidden_states, attn_metadata,
385
- # residual, mamba_cache_params, and sequence_idx.
373
+ # residual, and sequence_idx.
374
ssm_hidden, _ = self.mamba(
375
hidden_states=hidden_states * self.ssm_in_multiplier,
376
residual=residual,
389
- mamba_cache_params=mamba_cache_params,
390
- mamba2_metadata=mamba2_metadata,
377
**kwargs,
378
)
379
# Sum the outputs from both branches.
@@ -464,25 +450,10 @@ class FalconH1Model(nn.Module):
450
self,
451
input_ids: torch.Tensor,
452
positions: torch.Tensor,
467
- mamba_cache_params: MambaCacheParams,
453
intermediate_tensors: Optional[IntermediateTensors] = None,
454
inputs_embeds: Optional[torch.Tensor] = None,
455
) -> torch.Tensor:
456
472
- # pass a sequence index tensor, that is required for
473
- # proper continuous batching computation including
474
- # chunked prefill
475
- attn_metadata = get_forward_context().attn_metadata
476
-
477
- if not envs.VLLM_USE_V1:
478
- mamba2_metadata = prepare_mamba2_metadata(
479
- chunk_size=self.config.mamba_chunk_size,
480
- attn_metadata=attn_metadata,
481
- )
482
- else:
483
- # v1 get mamba2_metadata from forward_context
484
- mamba2_metadata = None
485
-
457
if get_pp_group().is_first_rank:
458
if inputs_embeds is not None:
459
hidden_states = inputs_embeds * self.embedding_multiplier
@@ -495,14 +466,9 @@ class FalconH1Model(nn.Module):
466
467
for i in range(self.start_layer, self.end_layer):
468
layer = self.layers[i]
498
- layer_mamba_cache_params = None
499
- if mamba_cache_params:
500
- layer_mamba_cache_params = mamba_cache_params.at_layer_idx(i)
469
hidden_states = layer(
470
positions=positions,
471
hidden_states=hidden_states,
504
- mamba_cache_params=layer_mamba_cache_params,
505
- mamba2_metadata=mamba2_metadata,
472
)
473
if not get_pp_group().is_last_rank:
474
return IntermediateTensors({
@@ -541,13 +507,11 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
507
def get_mamba_state_shape_from_config(
508
cls,
509
vllm_config: "VllmConfig",
544
- use_v1: bool = True,
510
) -> tuple[tuple[int, int], tuple[int, int, int]]:
511
"""Calculate shapes for Mamba's convolutional and state caches.
512
513
Args:
514
vllm_config: vLLM config
550
- use_v1: Get shapes for V1 (or V0)
515
516
Returns:
517
Tuple containing:
@@ -570,7 +534,6 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
534
head_dim=hf_config.mamba_d_head,
535
state_size=hf_config.mamba_d_state,
536
conv_kernel=hf_config.mamba_d_conv,
573
- use_v1=use_v1,
537
)
538
539
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -592,7 +555,6 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
555
prefix=maybe_prefix(prefix, "model"))
556
self.tie_word_embeddings = config.tie_word_embeddings
557
self.unpadded_vocab_size = config.vocab_size
595
- self.mamba_cache: Optional[MambaCacheManager] = None
558
if lora_config:
559
self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
560
if get_pp_group().is_last_rank:
@@ -637,40 +599,15 @@ class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
599
**kwargs,
600
):
601
640
- mamba_cache_params = None
641
- if not envs.VLLM_USE_V1:
642
- if self.mamba_cache is None:
643
- mamba_state_shape = \
644
- self.get_mamba_state_shape_from_config(
645
- self.vllm_config, use_v1=False)
646
- mamba_state_dtype = \
647
- self.get_mamba_state_dtype_from_config(
648
- self.vllm_config)
649
- self.mamba_cache = MambaCacheManager(
650
- self.vllm_config,
651
- self.config.num_hidden_layers,
652
- *mamba_state_shape,
653
- *mamba_state_dtype,
654
- )
655
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
656
-
602
hidden_states = self.model(
603
input_ids,
604
positions,
660
- mamba_cache_params,
605
intermediate_tensors,
606
inputs_embeds,
607
)
608
609
return hidden_states
610
667
- def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
668
- return self.mamba_cache.copy_inputs_before_cuda_graphs(
669
- input_buffers, **kwargs)
670
-
671
- def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
672
- return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
673
-
611
def compute_logits(
612
self,
613
hidden_states: torch.Tensor,
vllm/model_executor/models/granitemoehybrid.py
+6
-73
@@ -9,19 +9,15 @@ import torch
9
from torch import nn
10
from transformers import GraniteMoeHybridConfig
11
12
-from vllm import envs
12
from vllm.attention.layer import Attention
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import CacheConfig, ModelConfig, VllmConfig
15
from vllm.distributed import get_tensor_model_parallel_world_size
16
from vllm.distributed.parallel_state import get_pp_group
18
-from vllm.forward_context import get_forward_context
17
from vllm.model_executor.layers.layernorm import RMSNorm
18
from vllm.model_executor.layers.linear import (QKVParallelLinear,
19
RowParallelLinear)
20
from vllm.model_executor.layers.logits_processor import LogitsProcessor
23
-from vllm.model_executor.layers.mamba.mamba2_metadata import (
24
- Mamba2Metadata, prepare_mamba2_metadata)
21
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
22
from vllm.model_executor.layers.mamba.mamba_utils import (
23
MambaStateDtypeCalculator, MambaStateShapeCalculator)
@@ -30,10 +26,7 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
26
from vllm.model_executor.layers.vocab_parallel_embedding import (
27
DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
28
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
33
-from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
34
- MambaCacheParams)
29
from vllm.sequence import IntermediateTensors
36
-from vllm.utils import LayerBlockType
30
31
from .granitemoe import GraniteMoeMoE
32
from .granitemoeshared import GraniteMoeSharedMLP
@@ -102,14 +95,12 @@ class GraniteMoeHybridMambaDecoderLayer(nn.Module):
95
self,
96
hidden_states: torch.Tensor,
97
residual: Optional[torch.Tensor],
105
- mamba_cache_params: MambaCacheParams,
106
- mamba2_metadata: Mamba2Metadata,
98
**kwargs,
99
):
100
residual = hidden_states
101
hidden_states = self.input_layernorm(hidden_states)
102
output = torch.empty_like(hidden_states)
112
- self.mamba(hidden_states, output, mamba_cache_params, mamba2_metadata)
103
+ self.mamba(hidden_states, output)
104
hidden_states = residual + output * self.residual_multiplier
105
106
residual = hidden_states
@@ -182,8 +173,6 @@ class GraniteMoeHybridAttentionDecoderLayer(nn.Module):
173
positions: torch.Tensor,
174
hidden_states: torch.Tensor,
175
residual: Optional[torch.Tensor],
185
- mamba_cache_params: MambaCacheParams,
186
- mamba2_metadata: Mamba2Metadata,
176
) -> torch.Tensor:
177
residual = hidden_states
178
hidden_states = self.input_layernorm(hidden_states)
@@ -366,22 +355,10 @@ class GraniteMoeHybridModel(nn.Module):
355
self,
356
input_ids: torch.Tensor,
357
positions: torch.Tensor,
369
- mamba_cache_params: MambaCacheParams,
358
intermediate_tensors: Optional[IntermediateTensors] = None,
359
inputs_embeds: Optional[torch.Tensor] = None,
360
) -> torch.Tensor:
361
374
- attn_metadata = get_forward_context().attn_metadata
375
-
376
- if not envs.VLLM_USE_V1:
377
- mamba2_metadata = prepare_mamba2_metadata(
378
- chunk_size=self.config.mamba_chunk_size,
379
- attn_metadata=attn_metadata,
380
- )
381
- else:
382
- # v1 get mamba2_metadata from forward_context
383
- mamba2_metadata = None
384
-
362
if get_pp_group().is_first_rank:
363
if inputs_embeds is not None:
364
hidden_states = inputs_embeds
@@ -399,20 +376,9 @@ class GraniteMoeHybridModel(nn.Module):
376
for i, layer in enumerate(self.layers):
377
if isinstance(layer, GraniteMoeHybridAttentionDecoderLayer):
378
num_attn += 1
402
-
403
- layer_mamba_cache_params = None
404
- if isinstance(
405
- layer,
406
- GraniteMoeHybridMambaDecoderLayer) and mamba_cache_params:
407
- layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
408
- i - num_attn)
409
-
410
- hidden_states, residual = layer(
411
- positions=positions,
412
- hidden_states=hidden_states,
413
- residual=residual,
414
- mamba_cache_params=layer_mamba_cache_params,
415
- mamba2_metadata=mamba2_metadata)
379
+ hidden_states, residual = layer(positions=positions,
380
+ hidden_states=hidden_states,
381
+ residual=residual)
382
383
if not get_pp_group().is_last_rank:
384
return IntermediateTensors({
@@ -552,13 +518,11 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
518
def get_mamba_state_shape_from_config(
519
cls,
520
vllm_config: "VllmConfig",
555
- use_v1: bool = True,
521
) -> tuple[tuple[int, int], tuple[int, int, int]]:
522
"""Calculate shapes for Mamba's convolutional and state caches.
523
524
Args:
525
vllm_config: vLLM config
561
- use_v1: Get shapes for V1 (or V0)
526
527
Returns:
528
Tuple containing:
@@ -577,7 +541,6 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
541
head_dim=hf_config.mamba_d_head,
542
state_size=hf_config.mamba_d_state,
543
conv_kernel=hf_config.mamba_d_conv,
580
- use_v1=use_v1,
544
)
545
546
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -620,9 +583,6 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
583
scale=1 /
584
self.config.logits_scaling)
585
623
- # Used to track and store by the Mamba cache between steps.
624
- self.mamba_cache: Optional[MambaCacheManager] = None
625
-
586
self.make_empty_intermediate_tensors = (
587
self.model.make_empty_intermediate_tensors)
588
@@ -636,38 +596,11 @@ class GraniteMoeHybridForCausalLM(nn.Module, HasInnerState, SupportsLoRA,
596
inputs_embeds: Optional[torch.Tensor] = None,
597
**kwargs):
598
639
- mamba_cache_params = None
640
- if not envs.VLLM_USE_V1:
641
- if self.mamba_cache is None:
642
- num_mamba_layers = (
643
- self.model_config.get_num_layers_by_block_type(
644
- self.vllm_config.parallel_config,
645
- LayerBlockType.mamba))
646
- mamba_state_shape = \
647
- self.get_mamba_state_shape_from_config(
648
- self.vllm_config, use_v1=False)
649
- mamba_state_dtype = \
650
- self.get_mamba_state_dtype_from_config(
651
- self.vllm_config)
652
- self.mamba_cache = MambaCacheManager(self.vllm_config,
653
- num_mamba_layers,
654
- *mamba_state_shape,
655
- *mamba_state_dtype)
656
-
657
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
658
-
659
- hidden_states = self.model(input_ids, positions, mamba_cache_params,
660
- intermediate_tensors, inputs_embeds)
599
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
600
+ inputs_embeds)
601
602
return hidden_states
603
664
- def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
665
- return self.mamba_cache.copy_inputs_before_cuda_graphs(
666
- input_buffers, **kwargs)
667
-
668
- def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
669
- return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
670
-
604
def compute_logits(
605
self,
606
hidden_states: torch.Tensor,
vllm/model_executor/models/jamba.py
+8
-45
@@ -9,7 +9,6 @@ import torch
9
from torch import nn
10
from transformers import JambaConfig
11
12
-from vllm import envs
12
from vllm.attention.layer import Attention
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import CacheConfig, ModelConfig, VllmConfig
@@ -30,10 +29,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
29
DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
30
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
31
from vllm.model_executor.models.llama import LlamaMLP as JambaMLP
33
-from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
34
- MambaCacheParams)
32
from vllm.sequence import IntermediateTensors
36
-from vllm.utils import LayerBlockType
33
34
from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP
35
from .utils import (AutoWeightsLoader, WeightsMapper, is_pp_missing_parameter,
@@ -145,7 +141,6 @@ class JambaMambaDecoderLayer(nn.Module):
141
self,
142
hidden_states: torch.Tensor,
143
residual: Optional[torch.Tensor],
148
- mamba_cache_params: MambaCacheParams,
144
**kwargs,
145
):
146
if residual is None:
@@ -156,7 +151,7 @@ class JambaMambaDecoderLayer(nn.Module):
151
hidden_states, residual)
152
153
output = torch.empty_like(hidden_states)
159
- self.mamba(hidden_states, output, mamba_cache_params)
154
+ self.mamba(hidden_states, output)
155
# Fully Connected
156
hidden_states, residual = self.pre_ff_layernorm(output, residual)
157
hidden_states = self.feed_forward(hidden_states)
@@ -333,7 +328,6 @@ class JambaModel(nn.Module):
328
self,
329
input_ids: torch.Tensor,
330
positions: torch.Tensor,
336
- mamba_cache_params: MambaCacheParams,
331
intermediate_tensors: Optional[IntermediateTensors] = None,
332
inputs_embeds: Optional[torch.Tensor] = None,
333
) -> torch.Tensor:
@@ -348,24 +342,11 @@ class JambaModel(nn.Module):
342
hidden_states = intermediate_tensors["hidden_states"]
343
residual = intermediate_tensors["residual"]
344
351
- kv_cache_index = 0
352
- mamba_cache_index = 0
345
for layer in islice(self.layers, self.start_layer, self.end_layer):
354
- layer_mamba_cache_params = None
355
- if isinstance(layer, JambaAttentionDecoderLayer):
356
- kv_cache_index += 1
357
- if isinstance(layer,
358
- JambaMambaDecoderLayer) and mamba_cache_params:
359
- current_state_layer = mamba_cache_index
360
- layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
361
- current_state_layer)
362
- mamba_cache_index += 1
363
-
364
- hidden_states, residual = layer(
365
- positions=positions,
366
- hidden_states=hidden_states,
367
- residual=residual,
368
- mamba_cache_params=layer_mamba_cache_params)
346
+ hidden_states, residual = layer(positions=positions,
347
+ hidden_states=hidden_states,
348
+ residual=residual)
349
+
350
if not get_pp_group().is_last_rank:
351
return IntermediateTensors({
352
"hidden_states": hidden_states,
@@ -503,8 +484,6 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
484
if not lora_config else lora_config.lora_vocab_padding_size,
485
prefix=maybe_prefix(prefix, "lm_head"),
486
)
506
- # Used to track and store by the Mamba cache between steps.
507
- self.mamba_cache: Optional[MambaCacheManager] = None
487
488
self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
489
config.vocab_size)
@@ -521,24 +500,9 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
500
intermediate_tensors: Optional[IntermediateTensors] = None,
501
inputs_embeds: Optional[torch.Tensor] = None,
502
**kwargs):
524
- # NOTE: mamba_cache_params is not needed for v1
525
- mamba_cache_params = None
526
- if not envs.VLLM_USE_V1:
527
- if self.mamba_cache is None:
528
- num_layers = self.model_config.get_num_layers_by_block_type(
529
- self.vllm_config.parallel_config, LayerBlockType.mamba)
530
- state_shape = self.get_mamba_state_shape_from_config(
531
- self.vllm_config)
532
- state_dtype = self.get_mamba_state_dtype_from_config(
533
- self.vllm_config)
534
- self.mamba_cache = MambaCacheManager(self.vllm_config,
535
- num_layers, *state_shape,
536
- *state_dtype)
537
-
538
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
539
-
540
- hidden_states = self.model(input_ids, positions, mamba_cache_params,
541
- intermediate_tensors, inputs_embeds)
503
+
504
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
505
+ inputs_embeds)
506
return hidden_states
507
508
def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
@@ -574,7 +538,6 @@ class JambaForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
538
intermediate_size=hf_config.mamba_expand * hidden_size,
539
state_size=hf_config.mamba_d_state,
540
conv_kernel=hf_config.mamba_d_conv,
577
- use_v1=envs.VLLM_USE_V1,
541
)
542
543
def compute_logits(
vllm/model_executor/models/lfm2.py
-7
@@ -8,7 +8,6 @@ import torch
8
import torch.nn as nn
9
from transformers import Lfm2Config
10
11
-from vllm import envs
11
from vllm.attention import Attention
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, ModelConfig, VllmConfig
@@ -297,7 +296,6 @@ class Lfm2ShortConvDecoderLayer(nn.Module):
296
self.conv(
297
hidden_states,
298
output,
300
- conv_metadata=None,
299
)
300
hidden_states, residual = self.ffn_norm(output, residual)
301
hidden_states = self.feed_forward(hidden_states)
@@ -459,13 +457,11 @@ class Lfm2ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
457
def get_mamba_state_shape_from_config(
458
cls,
459
vllm_config: "VllmConfig",
462
- use_v1: bool = True,
460
) -> tuple[tuple[int, int]]:
461
""" Calculate shapes for LFM2's convolutional cache.
462
463
Args:
464
vllm_config: vLLM config
468
- use_v1: Get shapes for V1 (or V0)
465
466
Returns:
467
Tuple containing:
@@ -478,7 +474,6 @@ class Lfm2ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
474
tp_world_size=parallel_config.tensor_parallel_size,
475
intermediate_size=hf_config.conv_dim,
476
conv_kernel=hf_config.conv_L_cache,
481
- use_v1=use_v1,
477
)
478
479
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
@@ -489,8 +484,6 @@ class Lfm2ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
484
scheduler_config = vllm_config.scheduler_config
485
assert (not cache_config.enable_prefix_caching
486
), "Lfm2 currently does not support prefix caching"
492
- assert envs.VLLM_USE_V1, (
493
- "Lfm2ForCausalLM doesn't support vLLM v0. Please enable v1")
487
488
super().__init__()
489
self.config = config
vllm/model_executor/models/mamba.py
+6
-39
@@ -8,7 +8,6 @@ import torch
8
from torch import nn
9
from transformers import MambaConfig
10
11
-from vllm import envs
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import CacheConfig, ModelConfig, VllmConfig
13
from vllm.distributed.parallel_state import get_pp_group
@@ -24,10 +23,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
23
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
24
from vllm.model_executor.models.interfaces import (HasInnerState,
25
IsAttentionFree, SupportsPP)
27
-from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
28
- MambaCacheParams)
26
from vllm.sequence import IntermediateTensors
30
-from vllm.utils import LayerBlockType
27
28
from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
29
make_empty_intermediate_tensors_factory, make_layers,
@@ -72,7 +68,6 @@ class MambaDecoderLayer(nn.Module):
68
self,
69
hidden_states: torch.Tensor,
70
residual: Optional[torch.Tensor],
75
- mamba_cache_params: MambaCacheParams,
71
**kwargs,
72
):
73
if residual is None:
@@ -82,7 +77,7 @@ class MambaDecoderLayer(nn.Module):
77
hidden_states, residual = self.norm(hidden_states, residual)
78
79
output = torch.empty_like(hidden_states)
85
- self.mixer(hidden_states, output, mamba_cache_params)
80
+ self.mixer(hidden_states, output)
81
return output, residual
82
83
@@ -134,7 +129,6 @@ class MambaModel(nn.Module):
129
self,
130
input_ids: torch.Tensor,
131
positions: torch.Tensor,
137
- mamba_cache_params: Optional[MambaCacheParams] = None,
132
intermediate_tensors: Optional[IntermediateTensors] = None,
133
inputs_embeds: Optional[torch.Tensor] = None,
134
) -> torch.Tensor:
@@ -151,17 +145,9 @@ class MambaModel(nn.Module):
145
146
for i in range(self.start_layer, self.end_layer):
147
layer = self.layers[i]
154
-
155
- layer_cache_params = None
156
- if mamba_cache_params is not None:
157
- layer_cache_params = mamba_cache_params.at_layer_idx(
158
- i - self.start_layer)
159
-
160
- hidden_states, residual = layer(
161
- positions=positions,
162
- hidden_states=hidden_states,
163
- residual=residual,
164
- mamba_cache_params=layer_cache_params)
148
+ hidden_states, residual = layer(positions=positions,
149
+ hidden_states=hidden_states,
150
+ residual=residual)
151
if not get_pp_group().is_last_rank:
152
return IntermediateTensors({
153
"hidden_states": hidden_states,
@@ -225,9 +211,6 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):
211
prefix=maybe_prefix(prefix, "lm_head"),
212
)
213
228
- # Used to track and store by the Mamba cache between steps.
229
- self.mamba_cache: Optional[MambaCacheManager] = None
230
-
214
self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
215
config.vocab_size)
216
@@ -244,22 +227,7 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):
227
inputs_embeds: Optional[torch.Tensor] = None,
228
**kwargs):
229
247
- mamba_cache_params = None
248
- if not envs.VLLM_USE_V1:
249
- if self.mamba_cache is None:
250
- num_layers = self.model_config.get_num_layers_by_block_type(
251
- self.vllm_config.parallel_config, LayerBlockType.mamba)
252
- state_shape = self.get_mamba_state_shape_from_config(
253
- self.vllm_config)
254
- state_dtype = self.get_mamba_state_dtype_from_config(
255
- self.vllm_config)
256
- self.mamba_cache = MambaCacheManager(self.vllm_config,
257
- num_layers, *state_shape,
258
- *state_dtype)
259
-
260
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
261
-
262
- hidden_states = self.backbone(input_ids, positions, mamba_cache_params,
230
+ hidden_states = self.backbone(input_ids, positions,
231
intermediate_tensors, inputs_embeds)
232
233
return hidden_states
@@ -288,8 +256,7 @@ class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):
256
tp_world_size=parallel_config.tensor_parallel_size,
257
intermediate_size=hf_config.intermediate_size,
258
state_size=hf_config.state_size,
291
- conv_kernel=hf_config.conv_kernel,
292
- use_v1=envs.VLLM_USE_V1)
259
+ conv_kernel=hf_config.conv_kernel)
260
261
def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
262
return self.mamba_cache.copy_inputs_before_cuda_graphs(
vllm/model_executor/models/mamba2.py
+5
-58
@@ -8,16 +8,11 @@ import torch
8
from torch import nn
9
from transformers import MambaConfig
10
11
-from vllm import envs
12
-from vllm.attention.backends.abstract import AttentionMetadata
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import CacheConfig, ModelConfig, VllmConfig
13
from vllm.distributed.parallel_state import get_pp_group
16
-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
19
-from vllm.model_executor.layers.mamba.mamba2_metadata import (
20
- Mamba2Metadata, prepare_mamba2_metadata)
16
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
17
from vllm.model_executor.layers.mamba.mamba_utils import (
18
MambaStateDtypeCalculator, MambaStateShapeCalculator)
@@ -28,10 +23,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
23
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
24
from vllm.model_executor.models.interfaces import (HasInnerState,
25
IsAttentionFree)
31
-from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
32
- MambaCacheParams)
26
from vllm.sequence import IntermediateTensors
34
-from vllm.utils import LayerBlockType
27
28
from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
29
make_empty_intermediate_tensors_factory, make_layers,
@@ -74,8 +66,6 @@ class Mamba2DecoderLayer(nn.Module):
66
self,
67
hidden_states: torch.Tensor,
68
residual: Optional[torch.Tensor],
77
- mamba_cache_params: MambaCacheParams,
78
- mamba2_metadata: Mamba2Metadata,
69
**kwargs,
70
):
71
if residual is None:
@@ -85,7 +75,7 @@ class Mamba2DecoderLayer(nn.Module):
75
hidden_states, residual = self.norm(hidden_states, residual)
76
77
output = torch.empty_like(hidden_states)
88
- self.mixer(hidden_states, output, mamba_cache_params, mamba2_metadata)
78
+ self.mixer(hidden_states, output)
79
return output, residual
80
81
@@ -137,7 +127,6 @@ class Mamba2Model(nn.Module):
127
self,
128
input_ids: torch.Tensor,
129
positions: torch.Tensor,
140
- mamba_cache_params: MambaCacheParams,
130
intermediate_tensors: Optional[IntermediateTensors] = None,
131
inputs_embeds: Optional[torch.Tensor] = None,
132
) -> torch.Tensor:
@@ -152,25 +141,10 @@ class Mamba2Model(nn.Module):
141
hidden_states = intermediate_tensors["hidden_states"]
142
residual = intermediate_tensors["residual"]
143
155
- attn_metadata: AttentionMetadata = get_forward_context().attn_metadata
156
-
157
- if not envs.VLLM_USE_V1:
158
- mamba2_metadata = prepare_mamba2_metadata(
159
- chunk_size=self.config.chunk_size,
160
- attn_metadata=attn_metadata,
161
- )
162
- else:
163
- # v1 get mamba2_metadata from forward_context
164
- mamba2_metadata = None
165
-
144
for i, layer in enumerate(self.layers):
167
- hidden_states, residual = layer(
168
- positions=positions,
169
- hidden_states=hidden_states,
170
- residual=residual,
171
- mamba_cache_params=mamba_cache_params.at_layer_idx(
172
- i - self.start_layer) if mamba_cache_params else None,
173
- mamba2_metadata=mamba2_metadata)
145
+ hidden_states, residual = layer(positions=positions,
146
+ hidden_states=hidden_states,
147
+ residual=residual)
148
149
if not get_pp_group().is_last_rank:
150
return IntermediateTensors({
@@ -222,13 +196,11 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
196
def get_mamba_state_shape_from_config(
197
cls,
198
vllm_config: "VllmConfig",
225
- use_v1: bool = True,
199
) -> tuple[tuple[int, int], tuple[int, int, int]]:
200
"""Calculate shapes for Mamba's convolutional and state caches.
201
202
Args:
203
vllm_config: vLLM config
231
- use_v1: Get shapes for V1 (or V0)
204
205
Returns:
206
Tuple containing:
@@ -247,7 +219,6 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
219
head_dim=hf_config.head_dim,
220
state_size=hf_config.state_size,
221
conv_kernel=hf_config.conv_kernel,
250
- use_v1=use_v1,
222
)
223
224
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -282,9 +253,6 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
253
if config.tie_word_embeddings:
254
self.lm_head = self.lm_head.tie_weights(self.backbone.embeddings)
255
285
- # Used to track and store by the Mamba cache between steps.
286
- self.mamba_cache: Optional[MambaCacheManager] = None
287
-
256
self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
257
config.vocab_size)
258
@@ -300,29 +268,8 @@ class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
268
intermediate_tensors: Optional[IntermediateTensors] = None,
269
inputs_embeds: Optional[torch.Tensor] = None,
270
**kwargs):
303
- if not envs.VLLM_USE_V1:
304
- if self.mamba_cache is None:
305
- num_mamba_layers = (
306
- self.model_config.get_num_layers_by_block_type(
307
- self.vllm_config.parallel_config,
308
- LayerBlockType.mamba))
309
- mamba_state_shape = \
310
- self.get_mamba_state_shape_from_config(
311
- self.vllm_config, use_v1=False)
312
- mamba_state_dtype = \
313
- self.get_mamba_state_dtype_from_config(
314
- self.vllm_config)
315
- self.mamba_cache = MambaCacheManager(self.vllm_config,
316
- num_mamba_layers,
317
- *mamba_state_shape,
318
- *mamba_state_dtype)
319
-
320
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
321
- else:
322
- # NOTE: mamba_cache_params is not needed for v1
323
- mamba_cache_params = None
271
325
- hidden_states = self.backbone(input_ids, positions, mamba_cache_params,
272
+ hidden_states = self.backbone(input_ids, positions,
273
intermediate_tensors, inputs_embeds)
274
275
return hidden_states
vllm/model_executor/models/mamba_cache.py
deleted
-83
@@ -1,83 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from dataclasses import dataclass
5
-
6
-import torch
7
-
8
-from vllm.attention.backends.utils import PAD_SLOT_ID
9
-from vllm.config import VllmConfig
10
-from vllm.model_executor.models.constant_size_cache import ConstantSizeCache
11
-
12
-
13
-@dataclass
14
-class MambaCacheParams:
15
- conv_state: torch.Tensor = torch.Tensor()
16
- ssm_state: torch.Tensor = torch.Tensor()
17
- state_indices_tensor: torch.Tensor = torch.Tensor()
18
-
19
- def at_layer_idx(self, layer_idx):
20
- return MambaCacheParams(self.conv_state[layer_idx],
21
- self.ssm_state[layer_idx],
22
- self.state_indices_tensor)
23
-
24
-
25
-class MambaCacheManager(ConstantSizeCache):
26
-
27
- def __init__(self, vllm_config: VllmConfig, num_mamba_layers: int,
28
- conv_state_shape: tuple[int, int],
29
- temporal_state_shape: tuple[int, int],
30
- conv_state_dtype: torch.dtype,
31
- temporal_state_dtype: torch.dtype):
32
-
33
- self.conv_state_dtype = conv_state_dtype
34
- self.temporal_state_dtype = temporal_state_dtype
35
-
36
- # Determine max batch size to set size of MambaCache
37
- max_batch_size = vllm_config.scheduler_config.max_num_seqs
38
- if not vllm_config.model_config.enforce_eager:
39
- max_batch_size = vllm_config.pad_for_cudagraph(max_batch_size)
40
-
41
- # Initialize parent class
42
- super().__init__(max_batch_size)
43
-
44
- # assume conv_state = (dim, state_len)
45
- assert conv_state_shape[0] > conv_state_shape[1]
46
- conv_state = torch.empty(size=(num_mamba_layers, max_batch_size) +
47
- (conv_state_shape[1], conv_state_shape[0]),
48
- dtype=self.conv_state_dtype,
49
- device="cuda").transpose(-1, -2)
50
- temporal_state = torch.empty(size=(num_mamba_layers, max_batch_size) +
51
- temporal_state_shape,
52
- dtype=self.temporal_state_dtype,
53
- device="cuda")
54
-
55
- self._mamba_cache = (conv_state, temporal_state)
56
-
57
- @property
58
- def cache(self):
59
- return self._mamba_cache
60
-
61
- def _copy_cache(self, from_index: int, to_index: int):
62
- for cache_t in self.cache:
63
- cache_t[:, to_index].copy_(cache_t[:, from_index],
64
- non_blocking=True)
65
-
66
- def current_run_tensors(self, **kwargs) -> MambaCacheParams:
67
- """
68
- Return the tensors for the current run's conv and ssm state.
69
- """
70
- cache_tensors, state_indices_tensor = super().current_run_tensors(
71
- **kwargs)
72
- return MambaCacheParams(cache_tensors[0], cache_tensors[1],
73
- state_indices_tensor)
74
-
75
- def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
76
- """
77
- Provide the CUDA graph capture runs with a buffer in adjusted size.
78
- The buffer is used to maintain the Mamba Cache during the CUDA graph
79
- replay runs.
80
- """
81
- return self._mamba_cache, torch.as_tensor([PAD_SLOT_ID] * batch_size,
82
- dtype=torch.int32,
83
- device="cuda")
vllm/model_executor/models/minimax_cache.py
deleted
-36
@@ -1,36 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from dataclasses import dataclass
4
-
5
-import torch
6
-
7
-from vllm.model_executor.models.constant_size_cache import ConstantSizeCache
8
-
9
-
10
-@dataclass
11
-class MinimaxCacheParams:
12
- minimax_cache: torch.Tensor = torch.Tensor()
13
- state_indices_tensor: torch.Tensor = torch.Tensor()
14
-
15
- def at_layer_idx(self, layer_idx):
16
- return MinimaxCacheParams(self.minimax_cache[layer_idx, ...],
17
- self.state_indices_tensor)
18
-
19
-
20
-class MinimaxCacheManager(ConstantSizeCache):
21
-
22
- def __init__(self, dtype, cache_shape):
23
- super().__init__(cache_shape[1]) # max_batch_size is cache_shape[1]
24
- self._minimax_cache = torch.empty(size=cache_shape,
25
- dtype=dtype,
26
- device="cuda")
27
-
28
- @property
29
- def cache(self):
30
- return self._minimax_cache
31
-
32
- def _copy_cache(self, from_index: int, to_index: int):
33
- assert len(self.cache) > 0
34
- for cache_t in self.cache:
35
- cache_t[:, to_index].copy_(cache_t[:, from_index],
36
- non_blocking=True)
vllm/model_executor/models/minimax_text_01.py
-39
@@ -14,7 +14,6 @@ import torch.distributed
14
from torch import nn
15
from transformers import MiniMaxConfig
16
17
-from vllm import envs
17
from vllm.attention import Attention, AttentionMetadata
18
from vllm.compilation.decorators import support_torch_compile
19
from vllm.config import CacheConfig, ModelConfig, VllmConfig
@@ -44,7 +43,6 @@ from vllm.model_executor.models.utils import maybe_prefix
43
from vllm.sequence import IntermediateTensors
44
45
from .interfaces import HasInnerState, IsHybrid
47
-from .minimax_cache import MinimaxCacheManager, MinimaxCacheParams
46
from .utils import PPMissingLayer, is_pp_missing_parameter, make_layers
47
48
@@ -404,7 +402,6 @@ class MiniMaxText01DecoderLayer(nn.Module):
402
def forward(self,
403
hidden_states: torch.Tensor,
404
positions: torch.Tensor,
407
- kv_caches: Union[list[dict], Optional[torch.Tensor]],
405
attn_metadata: AttentionMetadata,
406
residual: Optional[torch.Tensor],
407
is_warmup: bool = False,
@@ -418,7 +415,6 @@ class MiniMaxText01DecoderLayer(nn.Module):
415
hidden_states=layernorm_output,
416
output=self_attention_output,
417
positions=positions,
421
- kv_caches=kv_caches,
418
)
419
420
residual = residual * self.layernorm_attention_alpha
@@ -563,10 +559,6 @@ class MiniMaxText01Model(nn.Module):
559
self._dtype = _dummy.dtype
560
del _dummy
561
566
- if not envs.VLLM_USE_V1:
567
- self.minimax_cache = MinimaxCacheManager(
568
- dtype=torch.float32, cache_shape=self.cache_shape)
569
-
562
norm_kwargs = {}
563
if hasattr(config, "rms_norm_eps"):
564
norm_kwargs["eps"] = config.rms_norm_eps
@@ -614,25 +606,6 @@ class MiniMaxText01Model(nn.Module):
606
**kwargs) -> Union[torch.Tensor, IntermediateTensors]:
607
forward_context = get_forward_context()
608
attn_metadata = forward_context.attn_metadata
617
- if not envs.VLLM_USE_V1 and attn_metadata is None:
618
- return None
619
- if not envs.VLLM_USE_V1:
620
- if "request_ids_to_seq_ids" not in kwargs:
621
- kwargs["request_ids_to_seq_ids"] = {}
622
- if "finished_requests_ids" not in kwargs:
623
- kwargs["finished_requests_ids"] = []
624
- (
625
- minimax_cache_tensors,
626
- state_indices_tensor,
627
- ) = self.minimax_cache.current_run_tensors(**kwargs)
628
- if getattr(attn_metadata, "num_prefills", 0) > 0:
629
- self._clear_prefill_cache(attn_metadata, minimax_cache_tensors,
630
- **kwargs)
631
-
632
- minimax_cache_params = MinimaxCacheParams(minimax_cache_tensors,
633
- state_indices_tensor)
634
- else:
635
- minimax_cache_params = None
609
610
if get_pp_group().is_first_rank:
611
if inputs_embeds is None:
@@ -645,20 +618,10 @@ class MiniMaxText01Model(nn.Module):
618
hidden_states = intermediate_tensors["hidden_states"]
619
residual = intermediate_tensors["residual"]
620
648
- minimax_cache_index = 0
649
-
621
for layer in islice(self.layers, self.start_layer, self.end_layer):
651
- _caches = None
652
- if not envs.VLLM_USE_V1 and isinstance(
653
- layer.self_attn, MiniMaxText01LinearAttention):
654
- current_state_layer = minimax_cache_index
655
- _caches = minimax_cache_params.at_layer_idx(
656
- current_state_layer)
657
- minimax_cache_index += 1
622
hidden_states, residual = layer(
623
hidden_states=hidden_states,
624
positions=positions,
661
- kv_caches=_caches,
625
attn_metadata=attn_metadata,
626
residual=residual,
627
)
@@ -1003,13 +966,11 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):
966
def get_mamba_state_shape_from_config(
967
cls,
968
vllm_config: "VllmConfig",
1006
- use_v1: bool = True,
969
) -> tuple[tuple[int, ...], ...]:
970
"""Calculate shape for MiniMaxText01LinearAttention cache.
971
972
Args:
973
vllm_config: vLLM config
1012
- use_v1: Get shapes for V1 (or V0)
974
975
Returns:
976
Tuple containing:
vllm/model_executor/models/nemotron_h.py
+3
-69
@@ -23,21 +23,17 @@ from typing import Optional
23
import torch
24
from torch import nn
25
26
-from vllm import envs
26
from vllm.attention.layer import Attention
27
from vllm.compilation.decorators import support_torch_compile
28
from vllm.config import CacheConfig, ModelConfig, VllmConfig
29
from vllm.distributed import get_tensor_model_parallel_world_size
30
from vllm.distributed.parallel_state import get_pp_group
32
-from vllm.forward_context import get_forward_context
31
from vllm.model_executor.layers.activation import ReLUSquaredActivation
32
from vllm.model_executor.layers.layernorm import RMSNorm
33
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
34
QKVParallelLinear,
35
RowParallelLinear)
36
from vllm.model_executor.layers.logits_processor import LogitsProcessor
39
-from vllm.model_executor.layers.mamba.mamba2_metadata import (
40
- Mamba2Metadata, prepare_mamba2_metadata)
37
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
38
from vllm.model_executor.layers.mamba.mamba_utils import (
39
MambaStateDtypeCalculator, MambaStateShapeCalculator)
@@ -49,14 +45,11 @@ from vllm.model_executor.model_loader.weight_utils import (
45
from vllm.model_executor.models.interfaces import (HasInnerState, IsHybrid,
46
SupportsLoRA, SupportsPP,
47
SupportsQuant)
52
-from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
53
- MambaCacheParams)
48
from vllm.model_executor.models.utils import (
49
AutoWeightsLoader, WeightsMapper, make_empty_intermediate_tensors_factory,
50
make_layers, maybe_prefix)
51
from vllm.sequence import IntermediateTensors
52
from vllm.transformers_utils.configs import NemotronHConfig
59
-from vllm.utils import LayerBlockType
53
54
55
class NemotronHMLP(nn.Module):
@@ -181,8 +174,6 @@ class NemotronHMambaDecoderLayer(nn.Module):
174
self,
175
hidden_states: torch.Tensor,
176
residual: Optional[torch.Tensor],
184
- mamba_cache_params: MambaCacheParams,
185
- mamba2_metadata: Mamba2Metadata,
177
**kwargs,
178
):
179
if residual is None:
@@ -192,7 +183,7 @@ class NemotronHMambaDecoderLayer(nn.Module):
183
hidden_states, residual = self.norm(hidden_states, residual)
184
185
output = torch.empty_like(hidden_states)
195
- self.mixer(hidden_states, output, mamba_cache_params, mamba2_metadata)
186
+ self.mixer(hidden_states, output)
187
return output, residual
188
189
@@ -370,22 +361,10 @@ class NemotronHModel(nn.Module):
361
self,
362
input_ids: torch.Tensor,
363
positions: torch.Tensor,
373
- mamba_cache_params: MambaCacheParams,
364
intermediate_tensors: Optional[IntermediateTensors] = None,
365
inputs_embeds: Optional[torch.Tensor] = None,
366
) -> torch.Tensor:
367
378
- attn_metadata = get_forward_context().attn_metadata
379
-
380
- if not envs.VLLM_USE_V1:
381
- mamba2_metadata = prepare_mamba2_metadata(
382
- chunk_size=self.config.chunk_size,
383
- attn_metadata=attn_metadata,
384
- )
385
- else:
386
- # v1 get mamba2_metadata from forward_context
387
- mamba2_metadata = None
388
-
368
if get_pp_group().is_first_rank:
369
if inputs_embeds is not None:
370
hidden_states = inputs_embeds
@@ -398,22 +377,11 @@ class NemotronHModel(nn.Module):
377
residual = intermediate_tensors["residual"]
378
379
residual = None
401
- num_non_mamba_layers = 0
380
for i, layer in enumerate(self.layers):
403
- layer_mamba_cache_params = None
404
- if isinstance(layer,
405
- NemotronHMambaDecoderLayer) and mamba_cache_params:
406
- layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
407
- i - num_non_mamba_layers)
408
- else:
409
- num_non_mamba_layers += 1
410
-
381
hidden_states, residual = layer(
382
positions=positions,
383
hidden_states=hidden_states,
384
residual=residual,
415
- mamba_cache_params=layer_mamba_cache_params,
416
- mamba2_metadata=mamba2_metadata,
385
)
386
387
if not get_pp_group().is_last_rank:
@@ -508,13 +476,11 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
476
def get_mamba_state_shape_from_config(
477
cls,
478
vllm_config: "VllmConfig",
511
- use_v1: bool = True,
479
) -> tuple[tuple[int, int], tuple[int, int, int]]:
480
"""Calculate shapes for Mamba's convolutional and state caches.
481
482
Args:
483
vllm_config: vLLM config
517
- use_v1: Get shapes for V1 (or V0)
484
485
Returns:
486
Tuple containing:
@@ -533,7 +499,6 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
499
head_dim=hf_config.mamba_head_dim,
500
state_size=hf_config.ssm_state_size,
501
conv_kernel=hf_config.conv_kernel,
536
- use_v1=use_v1,
502
)
503
504
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -566,8 +531,6 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
531
if not lora_config else lora_config.lora_vocab_padding_size,
532
prefix=maybe_prefix(prefix, "lm_head"),
533
)
569
- # Used to track and store by the Mamba cache between steps.
570
- self.mamba_cache: Optional[MambaCacheManager] = None
534
535
self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
536
config.vocab_size)
@@ -584,40 +547,11 @@ class NemotronHForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
547
inputs_embeds: Optional[torch.Tensor] = None,
548
**kwargs):
549
587
- mamba_cache_params = None
588
- if not envs.VLLM_USE_V1:
589
- if self.mamba_cache is None:
590
-
591
- num_mamba_layers = \
592
- self.model_config.get_num_layers_by_block_type(
593
- self.vllm_config.parallel_config,
594
- LayerBlockType.mamba
595
- )
596
- mamba_state_shape = \
597
- self.get_mamba_state_shape_from_config(
598
- self.vllm_config, use_v1=False)
599
- mamba_state_dtype = \
600
- self.get_mamba_state_dtype_from_config(
601
- self.vllm_config)
602
- self.mamba_cache = MambaCacheManager(self.vllm_config,
603
- num_mamba_layers,
604
- *mamba_state_shape,
605
- *mamba_state_dtype)
606
-
607
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
608
-
609
- hidden_states = self.model(input_ids, positions, mamba_cache_params,
610
- intermediate_tensors, inputs_embeds)
550
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
551
+ inputs_embeds)
552
553
return hidden_states
554
614
- def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
615
- return self.mamba_cache.copy_inputs_before_cuda_graphs(
616
- input_buffers, **kwargs)
617
-
618
- def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
619
- return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
620
-
555
def compute_logits(
556
self,
557
hidden_states: torch.Tensor,
vllm/model_executor/models/phi4flash.py
deleted
-731
@@ -1,731 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import math
4
-from collections.abc import Iterable
5
-from typing import Optional, Union
6
-
7
-import torch
8
-import torch.nn as nn
9
-from transformers.activations import ACT2FN
10
-
11
-import vllm.envs as envs
12
-from vllm.attention import Attention, AttentionMetadata, AttentionType
13
-from vllm.attention.selector import _Backend
14
-from vllm.config import CacheConfig, VllmConfig
15
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
16
-from vllm.forward_context import ForwardContext, get_forward_context
17
-from vllm.logger import init_logger
18
-from vllm.model_executor.layers.linear import (ColumnParallelLinear,
19
- MergedColumnParallelLinear,
20
- RowParallelLinear)
21
-from vllm.model_executor.layers.logits_processor import LogitsProcessor
22
-from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
23
- causal_conv1d_fn, causal_conv1d_update)
24
-from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
25
- selective_scan_fn, selective_state_update)
26
-from vllm.model_executor.layers.vocab_parallel_embedding import (
27
- DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
28
-from vllm.model_executor.models.interfaces import (HasInnerState, IsHybrid,
29
- SupportsV0Only)
30
-from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
31
- MambaCacheParams)
32
-from vllm.sequence import IntermediateTensors
33
-
34
-from .utils import make_layers, maybe_prefix
35
-
36
-logger = init_logger(__name__)
37
-
38
-
39
-class SwiGLUActivation(nn.Module):
40
-
41
- def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
42
- return x1 * nn.functional.silu(x2)
43
-
44
-
45
-class SambaYMLP(nn.Module):
46
- """Gated Linear Unit.
47
-
48
- Reference:
49
- Language Modeling with Gated Convolutional Networks.
50
- https://arxiv.org/pdf/1612.08083v3.pdf.
51
-
52
- """
53
-
54
- def __init__(self, config):
55
- super().__init__()
56
-
57
- self.config = config
58
- self.fc1 = nn.Linear(config.hidden_size,
59
- 2 * config.intermediate_size,
60
- bias=False)
61
- self.fc2 = nn.Linear(config.intermediate_size,
62
- config.hidden_size,
63
- bias=False)
64
-
65
- self.activation_fn = ACT2FN[config.hidden_act]
66
-
67
- def forward(self, hidden_states):
68
- y = self.fc1(hidden_states)
69
- gate, y = y.chunk(2, dim=-1)
70
- y = y * self.activation_fn(gate)
71
- return self.fc2(y)
72
-
73
-
74
-def get_virtual_engine():
75
- forward_context: ForwardContext = get_forward_context()
76
- return forward_context.virtual_engine
77
-
78
-
79
-class SambaYAttention(nn.Module):
80
-
81
- def __init__(self,
82
- config,
83
- layer_idx: Optional[int] = None,
84
- yoco_cross: bool = False,
85
- cache_config: Optional[CacheConfig] = None,
86
- prefix: str = ""):
87
- super().__init__()
88
- if layer_idx is None:
89
- logger.warning_once(
90
- f"Instantiating {self.__class__.__name__} without passing "
91
- "a `layer_idx` is not recommended and will lead to errors "
92
- "during the forward call if caching is used. Please make "
93
- "sure to provide a `layer_idx` when creating this class.")
94
- self.hidden_size = config.hidden_size
95
- self.num_heads = config.num_attention_heads
96
- self.head_dim = self.hidden_size // self.num_heads
97
- self.num_key_value_heads = config.num_key_value_heads
98
- self.yoco_cross = yoco_cross
99
-
100
- if (self.head_dim * self.num_heads) != self.hidden_size:
101
- raise ValueError("hidden_size must be divisible by num_heads "
102
- f"(got `hidden_size`: {self.hidden_size} and "
103
- f"`num_heads`: {self.num_heads}).")
104
-
105
- op_size = self.num_heads * self.head_dim + 2 * (
106
- self.num_key_value_heads * self.head_dim)
107
- self.out_proj = nn.Linear(self.num_heads * self.head_dim,
108
- self.hidden_size,
109
- bias=True)
110
- if yoco_cross:
111
- self.Wqkv = nn.Linear(self.hidden_size,
112
- self.num_heads * self.head_dim,
113
- bias=True)
114
- else:
115
- self.Wqkv = nn.Linear(self.hidden_size, op_size, bias=True)
116
-
117
- # disable sliding window for the second half of the model
118
- is_sliding = config.layer_types[layer_idx] == "sliding_attention"
119
- sliding_window = config.sliding_window if is_sliding else None
120
-
121
- assert self.num_heads % 2 == 0, 'num_heads should be even'
122
- assert self.num_key_value_heads % 2 == 0, 'num_heads should be even'
123
-
124
- self.lambda_init = self.lambda_init_fn(layer_idx)
125
- self.lambda_q1 = nn.Parameter(
126
- torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,
127
- std=0.1))
128
- self.lambda_k1 = nn.Parameter(
129
- torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,
130
- std=0.1))
131
- self.lambda_q2 = nn.Parameter(
132
- torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,
133
- std=0.1))
134
- self.lambda_k2 = nn.Parameter(
135
- torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,
136
- std=0.1))
137
- self.subln = nn.RMSNorm(2 * self.head_dim,
138
- eps=1e-5,
139
- elementwise_affine=True)
140
-
141
- params = {
142
- 'differential_flash_attention_config': {
143
- 'lambda_init': self.lambda_init,
144
- 'lambda_q1': self.lambda_q1,
145
- 'lambda_k1': self.lambda_k1,
146
- 'lambda_q2': self.lambda_q2,
147
- 'lambda_k2': self.lambda_k2,
148
- "subln": self.subln,
149
- }
150
- }
151
-
152
- if yoco_cross:
153
- kv_shared_layer_index = config.num_hidden_layers // 2 + 1
154
- kv_sharing_target_layer_name = \
155
- f"model.layers.{kv_shared_layer_index}.self_attn.attn"
156
- else:
157
- kv_sharing_target_layer_name = None
158
-
159
- self.attn = Attention(
160
- self.num_heads,
161
- self.head_dim,
162
- self.head_dim**-0.5,
163
- num_kv_heads=self.num_key_value_heads,
164
- cache_config=cache_config,
165
- per_layer_sliding_window=sliding_window,
166
- prefix=f"{prefix}.attn",
167
- attn_type=AttentionType.DECODER,
168
- kv_sharing_target_layer_name=kv_sharing_target_layer_name,
169
- **params)
170
- assert self.attn.backend == _Backend.DIFFERENTIAL_FLASH_ATTN,\
171
- "DIFFERENTIAL_FLASH_ATTN required"
172
-
173
- def lambda_init_fn(self, depth):
174
- return 0.8 - 0.6 * math.exp(-0.3 * depth)
175
-
176
- def forward(
177
- self,
178
- hidden_states: torch.Tensor,
179
- ):
180
-
181
- if not self.yoco_cross: # need to generate kv-cache
182
- qkv = self.Wqkv(hidden_states)
183
- q, k, v = qkv.split([
184
- self.hidden_size, self.num_key_value_heads * self.head_dim,
185
- self.num_key_value_heads * self.head_dim
186
- ],
187
- dim=-1)
188
- attn_output = self.attn(q, k, v)
189
- else: # reuse the kv cache, full attention
190
- q = self.Wqkv(hidden_states)
191
- attn_output = self.attn(q, None, None)
192
- attn_output = attn_output.view(-1, self.num_heads * self.head_dim)
193
- return self.out_proj(attn_output)
194
-
195
-
196
-class Phi4Mamba(nn.Module):
197
-
198
- def __init__(
199
- self,
200
- d_model,
201
- d_state=16,
202
- d_conv=4,
203
- expand=2,
204
- dt_rank="auto",
205
- dt_min=0.001,
206
- dt_max=0.1,
207
- dt_init="random", # difference
208
- dt_scale=1.0, # difference
209
- dt_init_floor=1e-4,
210
- conv_bias=True,
211
- bias=False,
212
- use_fast_path=True, # Fused kernel options
213
- layer_idx=None,
214
- device=None,
215
- dtype=None,
216
- yoco_cross=False,
217
- yoco_kv=False,
218
- ):
219
- factory_kwargs = {"params_dtype": dtype} # difference
220
- super().__init__()
221
- self.yoco_cross = yoco_cross
222
- self.yoco_kv = yoco_kv
223
- self.d_model = d_model
224
- self.d_state = d_state
225
- self.d_conv = d_conv
226
- self.expand = expand
227
- self.d_inner = int(self.expand * self.d_model)
228
- self.dt_rank = math.ceil(self.d_model /
229
- 16) if dt_rank == "auto" else dt_rank
230
- self.use_fast_path = use_fast_path
231
- self.layer_idx = layer_idx
232
- self.swiGluActivation = SwiGLUActivation()
233
- if self.yoco_cross:
234
- self.in_proj = MergedColumnParallelLinear(self.d_model,
235
- [self.d_inner],
236
- bias=bias,
237
- **factory_kwargs)
238
- self.out_proj = RowParallelLinear(self.d_inner,
239
- self.d_model,
240
- bias=bias,
241
- **factory_kwargs)
242
- return
243
- self.conv1d = ColumnParallelLinear(
244
- input_size=d_conv,
245
- output_size=self.d_inner,
246
- bias=conv_bias,
247
- params_dtype=dtype,
248
- )
249
- # unsqueeze to fit conv1d weights shape into the linear weights shape.
250
- # Can't do this in `weight_loader` since it already exists in
251
- # `ColumnParallelLinear` and `set_weight_attrs`
252
- # doesn't allow to override it
253
- self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
254
-
255
- self.in_proj = MergedColumnParallelLinear(
256
- self.d_model,
257
- [self.d_inner] * 2,
258
- bias=bias,
259
- params_dtype=dtype,
260
- )
261
-
262
- # selective projection used to make dt, B and C input dependent
263
- self.x_proj = RowParallelLinear(
264
- self.d_inner,
265
- self.dt_rank + self.d_state * 2,
266
- bias=False,
267
- params_dtype=dtype,
268
- )
269
-
270
- # time step projection (discretization) -
271
- # In the forward we need to apply dt_proj without the bias,
272
- # as the bias is added in the selective scan kernel.
273
- self.dt_proj = ColumnParallelLinear(
274
- self.dt_rank,
275
- self.d_inner,
276
- bias=True,
277
- skip_bias_add=True,
278
- params_dtype=dtype,
279
- )
280
-
281
- # # D "skip" parameter
282
- # self.D = nn.Parameter(torch.ones(self.d_inner)) # Keep in fp32
283
- self.A = nn.Parameter(
284
- torch.empty(
285
- self.d_inner,
286
- self.d_state,
287
- dtype=torch.float32,
288
- ))
289
- self.D = nn.Parameter(torch.ones(self.d_inner, dtype=torch.float32))
290
-
291
- self.out_proj = RowParallelLinear(
292
- self.d_inner,
293
- self.d_model,
294
- bias=bias,
295
- input_is_parallel=True,
296
- params_dtype=dtype,
297
- )
298
- self.activation = "silu"
299
-
300
- def forward(self,
301
- hidden_states: torch.Tensor,
302
- attn_metadata: AttentionMetadata,
303
- mamba_cache_params: MambaCacheParams,
304
- yoco_key_values=None) -> torch.Tensor:
305
-
306
- if self.yoco_cross:
307
- out = self.in_proj(hidden_states)[0]
308
- out = self.swiGluActivation(yoco_key_values, out)
309
- out = self.out_proj(out)
310
- return out[0], yoco_key_values
311
-
312
- # 1. Gated MLP's linear projection
313
- # projected_states = self.in_proj(hidden_states)[0].transpose(-2, -1)
314
- projected_states = self.in_proj(
315
- hidden_states.to(self.in_proj.weight.dtype))[0].transpose(-2, -1)
316
- hidden_states, gate = projected_states.chunk(2, dim=-2)
317
-
318
- # 2. Convolution sequence transformation
319
- conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),
320
- self.conv1d.weight.size(2))
321
-
322
- if attn_metadata.query_start_loc is not None \
323
- and attn_metadata.context_lens_tensor is not None:
324
- # |---------- N-1 iteration --------|
325
- # |---------------- N iteration ---------------------|
326
- # |- tokenA -|......................|-- newTokens ---|
327
- # |---------- context_len ----------|
328
- # |-------------------- seq_len ---------------------|
329
- # |-- query_len ---|
330
- hidden_states = causal_conv1d_fn(
331
- hidden_states,
332
- conv_weights,
333
- self.conv1d.bias,
334
- activation=self.activation,
335
- conv_states=mamba_cache_params.conv_state,
336
- has_initial_state=attn_metadata.context_lens_tensor > 0,
337
- cache_indices=mamba_cache_params.state_indices_tensor,
338
- query_start_loc=attn_metadata.query_start_loc)
339
- else:
340
- hidden_states = causal_conv1d_update(
341
- hidden_states.transpose(0, 1),
342
- mamba_cache_params.conv_state,
343
- conv_weights,
344
- self.conv1d.bias,
345
- self.activation,
346
- conv_state_indices=mamba_cache_params.state_indices_tensor)
347
- hidden_states = hidden_states.transpose(0, 1)
348
-
349
- # 3. State Space Model sequence transformation
350
- # 3.a. input varying initialization of time_step, B and C
351
- ssm_parameters = self.x_proj(hidden_states.transpose(-2, -1))[0]
352
-
353
- time_step, B, C = torch.split(
354
- ssm_parameters,
355
- [self.dt_rank, self.d_state, self.d_state],
356
- dim=-1,
357
- )
358
-
359
- # Note that Jamba normalizes B, C, and time_step here but Mamba doesn't.
360
-
361
- discrete_time_step = self.dt_proj(time_step)[0].transpose(-2, -1)
362
- # 3.c perform the recurrence y ← SSM(A, B, C)(x)
363
- time_proj_bias = (self.dt_proj.bias.float() if hasattr(
364
- self.dt_proj, "bias") else None)
365
-
366
- if attn_metadata.query_start_loc is not None \
367
- and attn_metadata.context_lens_tensor is not None:
368
- scan_outputs = selective_scan_fn(
369
- hidden_states,
370
- mamba_cache_params.ssm_state,
371
- discrete_time_step,
372
- self.A,
373
- B.transpose(-2, -1),
374
- C.transpose(-2, -1),
375
- self.D.float(),
376
- # z,
377
- None if self.yoco_kv else gate,
378
- time_proj_bias,
379
- delta_softplus=True,
380
- cache_indices=mamba_cache_params.state_indices_tensor,
381
- has_initial_state=attn_metadata.context_lens_tensor > 0,
382
- query_start_loc=attn_metadata.query_start_loc)
383
- else:
384
- scan_outputs = torch.empty_like(hidden_states.transpose(0, 1))
385
- selective_state_update(
386
- mamba_cache_params.ssm_state,
387
- hidden_states.transpose(0, 1),
388
- discrete_time_step.transpose(0, 1),
389
- self.A,
390
- B,
391
- C,
392
- self.D,
393
- # z
394
- # gate.transpose(0, 1),
395
- None if self.yoco_kv else gate.transpose(0, 1),
396
- time_proj_bias,
397
- dt_softplus=True,
398
- state_batch_indices=mamba_cache_params.state_indices_tensor,
399
- out=scan_outputs)
400
- scan_outputs = scan_outputs.transpose(0, 1)
401
-
402
- # 4. Final linear projection
403
- if self.yoco_kv:
404
- # gate = gate.transpose(-1,-2).contiguous()
405
- yoco_key_values = scan_outputs.transpose(-2, -1)
406
- scan_outputs = self.swiGluActivation(scan_outputs, gate)
407
-
408
- contextualized_states = self.out_proj(scan_outputs.transpose(-2,
409
- -1))[0]
410
-
411
- return contextualized_states, yoco_key_values
412
-
413
-
414
-class SambaYDecoderLayer(nn.Module):
415
-
416
- def __init__(
417
- self,
418
- config,
419
- layer_idx,
420
- cache_config,
421
- prefix: str = "",
422
- ) -> None:
423
- super().__init__()
424
-
425
- self.config = config
426
- self.layer_idx = layer_idx
427
-
428
- self.mlp = SambaYMLP(config)
429
- self.input_layernorm = nn.LayerNorm(config.hidden_size,
430
- eps=config.layer_norm_eps)
431
-
432
- self.yoco_mb = False
433
- self.yoco_cross = False
434
- if layer_idx >= config.num_hidden_layers // 2:
435
- self.yoco_mb = True
436
- self.yoco_cross = (layer_idx
437
- >= (config.num_hidden_layers // 2 + 2))
438
- self.use_mamba = config.mb_per_layer > 0 and \
439
- layer_idx % config.mb_per_layer == 0
440
- if self.use_mamba:
441
- factory_kwargs = {"dtype": None}
442
- self.attn = Phi4Mamba(config.hidden_size,
443
- layer_idx=layer_idx,
444
- yoco_cross=self.yoco_cross,
445
- yoco_kv=self.yoco_mb,
446
- **factory_kwargs)
447
- else:
448
- self.attn = SambaYAttention(config,
449
- layer_idx=layer_idx,
450
- yoco_cross=self.yoco_cross,
451
- cache_config=cache_config,
452
- prefix=f"{prefix}.self_attn")
453
- self.post_attention_layernorm = nn.LayerNorm(config.hidden_size,
454
- eps=config.layer_norm_eps)
455
-
456
- def forward(
457
- self,
458
- hidden_states: torch.Tensor,
459
- positions: torch.Tensor,
460
- attn_metadata: AttentionMetadata,
461
- mamba_cache_params: MambaCacheParams,
462
- ssm_output: Optional[torch.LongTensor] = None,
463
- ) -> Union[torch.Tensor, IntermediateTensors]:
464
- if self.use_mamba:
465
- assert mamba_cache_params is not None
466
- else:
467
- assert mamba_cache_params is None
468
-
469
- residual = hidden_states
470
- hidden_states = self.input_layernorm(
471
- hidden_states.to(dtype=self.input_layernorm.weight.dtype))
472
-
473
- if self.use_mamba:
474
- attn_outputs, ssm_output = self.attn(hidden_states,
475
- attn_metadata,
476
- mamba_cache_params,
477
- yoco_key_values=ssm_output)
478
- residual = residual.to(torch.float32)
479
- else:
480
- attn_outputs = self.attn(hidden_states, )
481
- hidden_states = residual + attn_outputs
482
- residual = hidden_states
483
- hidden_states = self.post_attention_layernorm(
484
- hidden_states.to(dtype=self.post_attention_layernorm.weight.dtype))
485
- hidden_states = self.mlp(hidden_states)
486
- hidden_states = residual + hidden_states
487
-
488
- return hidden_states, ssm_output
489
-
490
-
491
-class SambaYModel(nn.Module):
492
-
493
- def __init__(self,
494
- config,
495
- cache_config=None,
496
- quant_config=None,
497
- lora_config=None,
498
- prefix: str = "") -> None:
499
- super().__init__()
500
- self.config = config
501
- self.vocab_size = config.vocab_size
502
- self.embed_tokens = VocabParallelEmbedding(
503
- self.vocab_size,
504
- config.hidden_size,
505
- org_num_embeddings=config.vocab_size,
506
- )
507
-
508
- # Pipeline parallel is not supported since the second half of
509
- # the layers share the kv cache.
510
- if get_pp_group().world_size != 1:
511
- raise ValueError("Pipeline Parallel not supported")
512
-
513
- self.start_layer, self.end_layer, self.layers = make_layers(
514
- config.num_hidden_layers,
515
- lambda prefix: SambaYDecoderLayer(config,
516
- int(prefix.split('.')[-1]),
517
- cache_config,
518
- prefix=prefix),
519
- prefix=f"{prefix}.layers")
520
- self.final_layernorm = nn.LayerNorm(config.hidden_size,
521
- eps=config.layer_norm_eps)
522
-
523
- def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
524
- return self.embed_tokens(input_ids)
525
-
526
- def forward(
527
- self,
528
- input_ids: Optional[torch.Tensor],
529
- positions: torch.Tensor,
530
- attn_metadata: AttentionMetadata,
531
- mamba_cache_params: MambaCacheParams,
532
- intermediate_tensors: Optional[IntermediateTensors] = None,
533
- inputs_embeds: Optional[torch.Tensor] = None,
534
- ) -> Union[torch.Tensor, IntermediateTensors]:
535
-
536
- if get_pp_group().is_first_rank:
537
- if inputs_embeds is not None:
538
- hidden_states = inputs_embeds
539
- else:
540
- hidden_states = self.get_input_embeddings(input_ids)
541
- else:
542
- assert intermediate_tensors is not None
543
- hidden_states = intermediate_tensors["hidden_states"]
544
-
545
- mamba_state_idx = 0
546
- ssm_output = None
547
- for i in range(self.start_layer, self.end_layer):
548
- layer = self.layers[i]
549
- if i == self.config.num_hidden_layers // 2 + 2:
550
- # profile run
551
- kv_cache_idx = self.config.num_hidden_layers // 2 + 1
552
- cache_layer = self.layers[kv_cache_idx]
553
- kv_cache = cache_layer.attn.attn.kv_cache
554
- if kv_cache[0].numel() == 0:
555
- break
556
-
557
- # Starting from this layer, we do not need to calculate
558
- # the kv cache since we reuse the kv cache from last layer.
559
- # If in prefill phase, we can <s>prune></s> truncate
560
- # the hidden state to save computation cost.
561
- if attn_metadata.prefill_metadata and not envs.VLLM_USE_V1:
562
- selected_token_indices = torch.cumsum(
563
- attn_metadata.seq_lens_tensor, dim=0) - 1
564
- hidden_states = hidden_states.index_select(
565
- 0, selected_token_indices)
566
- ssm_output = ssm_output.index_select(
567
- 0, selected_token_indices)
568
-
569
- if layer.use_mamba:
570
- if i < self.config.num_hidden_layers // 2 or \
571
- not layer.yoco_cross:
572
- mamba_cache = mamba_cache_params.at_layer_idx(
573
- mamba_state_idx)
574
- mamba_state_idx += 1
575
- else:
576
- mamba_cache = mamba_cache_params.at_layer_idx(
577
- mamba_state_idx - 1)
578
-
579
- hidden_states, ssm_output = layer(hidden_states,
580
- positions,
581
- attn_metadata,
582
- mamba_cache,
583
- ssm_output=ssm_output)
584
- else:
585
- hidden_states, ssm_output = layer(
586
- hidden_states,
587
- positions,
588
- attn_metadata,
589
- None, # mamba_cache_params
590
- ssm_output=ssm_output)
591
-
592
- hidden_states = self.final_layernorm(
593
- hidden_states.to(dtype=self.final_layernorm.weight.dtype))
594
- return hidden_states
595
-
596
-
597
-class Phi4FlashForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsV0Only):
598
-
599
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
600
- config = vllm_config.model_config.hf_config
601
- cache_config = vllm_config.cache_config
602
- lora_config = vllm_config.lora_config
603
- quant_config = vllm_config.quant_config
604
- scheduler_config = vllm_config.scheduler_config
605
- self.compilation_config = vllm_config.compilation_config
606
- self.vllm_config = vllm_config
607
- # Prefix caching and chunked prefill is not supported for this model.
608
- assert not cache_config.enable_prefix_caching, \
609
- "Phi4flash currently does not support prefix caching"
610
- assert not scheduler_config.chunked_prefill_enabled, \
611
- "Phi4Flash currently does not support prefix caching"
612
- super().__init__()
613
- self.config = config
614
- self.model_config = vllm_config.model_config
615
- self.scheduler_config = scheduler_config
616
- self.model = SambaYModel(config,
617
- cache_config=cache_config,
618
- prefix=maybe_prefix(prefix, "model"))
619
- self.unpadded_vocab_size = config.vocab_size
620
- if lora_config:
621
- self.unpadded_vocab_size += lora_config.lora_extra_vocab_size
622
- self.lm_head = ParallelLMHead(
623
- self.unpadded_vocab_size,
624
- config.hidden_size,
625
- org_num_embeddings=config.vocab_size,
626
- padding_size=(
627
- DEFAULT_VOCAB_PADDING_SIZE
628
- # We need bigger padding if using lora for kernel
629
- # compatibility
630
- if not lora_config else lora_config.lora_vocab_padding_size),
631
- quant_config=quant_config,
632
- prefix=maybe_prefix(prefix, "lm_head"),
633
- )
634
- self.embedding_bias = None
635
- # Used to track and store by the Mamba cache between steps.
636
- self.mamba_cache: Optional[MambaCacheManager] = None
637
- self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
638
- config.vocab_size,
639
- logits_as_input=False)
640
-
641
- def forward(
642
- self,
643
- input_ids: torch.Tensor,
644
- positions: torch.Tensor,
645
- intermediate_tensors: Optional[IntermediateTensors] = None,
646
- inputs_embeds: Optional[torch.Tensor] = None,
647
- **kwargs,
648
- ) -> Union[torch.Tensor, IntermediateTensors]:
649
- if self.mamba_cache is None:
650
- num_mamba_layers = self.config.num_hidden_layers \
651
- // 2 // self.config.mb_per_layer + 1
652
- self.mamba_cache = MambaCacheManager(
653
- self.vllm_config,
654
- num_mamba_layers,
655
- *self._get_mamba_cache_shape(),
656
- self.lm_head.weight.dtype,
657
- self.lm_head.weight.dtype,
658
- )
659
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
660
-
661
- attn_metadata = get_forward_context().attn_metadata
662
- # input_ids and hidden_states isn't a one-to-one mapping in prefill
663
- # stage due to YOCO optimization.
664
- hidden_states = self.model(input_ids, positions, attn_metadata,
665
- mamba_cache_params, intermediate_tensors,
666
- inputs_embeds)
667
- return hidden_states
668
-
669
- def _get_mamba_cache_shape(
670
- self
671
- ) -> tuple[Optional[tuple[int, int]], Optional[tuple[int, int]]]:
672
- world_size = get_tensor_model_parallel_world_size()
673
- hidden_size = self.config.hidden_size
674
- mamba_expand = self.config.mamba_expand # 2
675
- mamba_d_conv = self.config.mamba_d_conv # 4
676
- mamba_d_state = self.config.mamba_d_state # 16
677
- conv_state_shape = (
678
- mamba_expand * hidden_size // world_size,
679
- mamba_d_conv - 1,
680
- )
681
- temporal_state_shape = (
682
- mamba_expand * hidden_size // world_size,
683
- mamba_d_state,
684
- )
685
- return conv_state_shape, temporal_state_shape
686
-
687
- def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
688
- return self.mamba_cache.copy_inputs_before_cuda_graphs(
689
- input_buffers, **kwargs)
690
-
691
- def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
692
- return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
693
-
694
- def compute_logits(
695
- self,
696
- hidden_states: torch.Tensor,
697
- ) -> Optional[torch.Tensor]:
698
- processed_logits = self.logits_processor(
699
- self.lm_head,
700
- hidden_states,
701
- self.embedding_bias,
702
- )
703
- return processed_logits
704
-
705
- def load_weights(
706
- self,
707
- weights: Iterable[tuple[str, torch.Tensor]],
708
- ):
709
- weights = {name: weight for name, weight in weights}
710
- adjusted_weights = {}
711
- for name, weight in weights.items():
712
- if "A_log" in name:
713
- name = name.replace("A_log", "A")
714
- weight = -torch.exp(weight.float())
715
- if "inner_cross_attn." in name:
716
- name = name.replace("inner_cross_attn.", "")
717
- adjusted_weights[name] = weight
718
- adjusted_weights["lm_head.weight"] = weights[
719
- "model.embed_tokens.weight"]
720
- loaded_params: set[str] = set()
721
- for name, param in self.named_parameters():
722
- weight = adjusted_weights.get(name)
723
- if weight is not None and weight.shape != param.shape:
724
- logger.warning("Shape mismatch: %s %s %s", name, weight.shape,
725
- param.shape)
726
- loaded_params.add(name)
727
- missing_keys, unexpected_keys = self.load_state_dict(adjusted_weights,
728
- strict=False)
729
- assert len(unexpected_keys) == 0, f"Unexpected keys: {unexpected_keys}"
730
- assert len(missing_keys) == 0, f"Missing keys: {missing_keys}"
731
- return loaded_params
vllm/model_executor/models/plamo2.py
+59
-181
@@ -12,7 +12,6 @@ import torch
12
from torch import nn
13
from transformers import PretrainedConfig
14
15
-from vllm import envs
15
from vllm.attention.backends.abstract import AttentionMetadata
16
from vllm.attention.layer import Attention
17
from vllm.compilation.decorators import support_torch_compile
@@ -29,8 +28,6 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
28
RowParallelLinear)
29
from vllm.model_executor.layers.logits_processor import LogitsProcessor
30
from vllm.model_executor.layers.mamba.abstract import MambaBase
32
-from vllm.model_executor.layers.mamba.mamba2_metadata import (
33
- Mamba2Metadata, prepare_mamba2_metadata, update_metadata)
31
from vllm.model_executor.layers.mamba.mamba_utils import (
32
MambaStateDtypeCalculator, MambaStateShapeCalculator)
33
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
@@ -47,15 +44,13 @@ from vllm.model_executor.model_loader.weight_utils import (
44
composed_weight_loader, default_weight_loader, sharded_weight_loader)
45
from vllm.model_executor.models.interfaces import (HasInnerState, IsHybrid,
46
SupportsPP)
50
-from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
51
- MambaCacheParams)
47
from vllm.model_executor.models.utils import (
48
is_pp_missing_parameter, make_empty_intermediate_tensors_factory,
49
make_layers, maybe_prefix)
50
from vllm.model_executor.utils import set_weight_attrs
51
from vllm.platforms import current_platform
52
from vllm.sequence import IntermediateTensors
58
-from vllm.utils import LayerBlockType, direct_register_custom_op
53
+from vllm.utils import direct_register_custom_op
54
from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
55
56
@@ -194,17 +189,13 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
189
190
self.chunk_size = self.config.mamba_chunk_size
191
197
- if envs.VLLM_USE_V1:
198
- compilation_config = get_current_vllm_config().compilation_config
199
- if prefix in compilation_config.static_forward_context:
200
- raise ValueError(f"Duplicate layer name: {prefix}")
201
- compilation_config.static_forward_context[prefix] = self
202
- # The outer list is for v0 PP virtual engine. Though this code path
203
- # only runs for v1, we have to do this to unify with the interface
204
- # of Attention + v0 PP.
205
- # The inner tuple is (conv_state, ssm_state)
206
- self.kv_cache = [(torch.tensor([]), torch.tensor([]))]
207
- assert self.chunk_size != -1, "chunk_size must be set for v1"
192
+ compilation_config = get_current_vllm_config().compilation_config
193
+ if prefix in compilation_config.static_forward_context:
194
+ raise ValueError(f"Duplicate layer name: {prefix}")
195
+ compilation_config.static_forward_context[prefix] = self
196
+ # The tuple is (conv_state, ssm_state)
197
+ self.kv_cache = (torch.tensor([]), torch.tensor([]))
198
+ assert self.chunk_size != -1, "chunk_size must be set for v1"
199
200
self.prefix = prefix
201
@@ -227,8 +218,6 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
218
self,
219
hidden_states: torch.Tensor,
220
output: torch.Tensor,
230
- mamba_cache_params: MambaCacheParams,
231
- mamba2_metadata: Mamba2Metadata,
221
**kwargs,
222
):
223
pass
@@ -237,59 +226,43 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
226
self,
227
hidden_states: torch.Tensor,
228
output: torch.Tensor,
240
- mamba_cache_params: MambaCacheParams,
241
- mamba2_metadata: Mamba2Metadata,
229
**kwargs,
230
):
244
- if not envs.VLLM_USE_V1:
245
- CustomOp.forward(self, hidden_states, output, mamba_cache_params,
246
- mamba2_metadata)
247
- else:
248
- torch.ops.vllm.plamo2_mamba_mixer(
249
- hidden_states,
250
- output,
251
- self.prefix,
252
- )
231
+ torch.ops.vllm.plamo2_mamba_mixer(
232
+ hidden_states,
233
+ output,
234
+ self.prefix,
235
+ )
236
237
def forward_cuda(
238
self,
239
hidden_states: torch.Tensor,
240
output: torch.Tensor,
258
- mamba_cache_params: MambaCacheParams,
259
- mamba2_metadata: Mamba2Metadata,
241
**kwargs,
242
):
243
244
forward_context = get_forward_context()
264
- # mamba2_metadata contains metadata necessary for the mamba2 triton
245
+ # attn_metadata contains metadata necessary for the mamba2 triton
246
# kernels to operate in continuous batching and in chunked prefill
247
# modes; they are computed at top-level model forward since they
248
# stay the same and reused for all mamba layers in the same iteration
249
attn_metadata: AttentionMetadata = forward_context.attn_metadata
269
- if envs.VLLM_USE_V1:
270
- if attn_metadata is not None:
271
- assert isinstance(attn_metadata, dict)
272
- attn_metadata = attn_metadata[self.prefix]
273
- mamba2_metadata = attn_metadata
274
- assert isinstance(attn_metadata, Mamba2AttentionMetadata)
275
- self_kv_cache = self.kv_cache[forward_context.virtual_engine]
276
- # conv_state = (..., dim, width-1) yet contiguous along 'dim'
277
- conv_state = self_kv_cache[0].transpose(-1, -2)
278
- ssm_state = self_kv_cache[1]
279
- state_indices_tensor = attn_metadata.state_indices_tensor
280
- else:
281
- conv_state = mamba_cache_params.conv_state
282
- ssm_state = mamba_cache_params.ssm_state
283
- state_indices_tensor = mamba_cache_params.state_indices_tensor
284
-
285
- # Common members between V1 metadata and V0 metadata
286
- if mamba2_metadata is not None:
287
- has_initial_states_p = mamba2_metadata.has_initial_states_p
288
- prep_initial_states = mamba2_metadata.prep_initial_states
289
- chunk_size = mamba2_metadata.chunk_size
290
- seq_idx_p = mamba2_metadata.seq_idx_p
291
- chunk_indices_p = mamba2_metadata.chunk_indices_p
292
- chunk_offsets_p = mamba2_metadata.chunk_offsets_p
250
+
251
+ if attn_metadata is not None:
252
+ assert isinstance(attn_metadata, dict)
253
+ attn_metadata = attn_metadata[self.prefix]
254
+ assert isinstance(attn_metadata, Mamba2AttentionMetadata)
255
+ self_kv_cache = self.kv_cache[forward_context.virtual_engine]
256
+ # conv_state = (..., dim, width-1) yet contiguous along 'dim'
257
+ conv_state = self_kv_cache[0].transpose(-1, -2)
258
+ ssm_state = self_kv_cache[1]
259
+ state_indices_tensor = attn_metadata.state_indices_tensor
260
+ has_initial_states_p = attn_metadata.has_initial_states_p
261
+ prep_initial_states = attn_metadata.prep_initial_states
262
+ chunk_size = attn_metadata.chunk_size
263
+ seq_idx_p = attn_metadata.seq_idx_p
264
+ chunk_indices_p = attn_metadata.chunk_indices_p
265
+ chunk_offsets_p = attn_metadata.chunk_offsets_p
266
267
# 1. Gated MLP's linear projection
268
projected_states = self.in_proj(hidden_states)
@@ -299,8 +272,8 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
272
conv_weights = self.conv1d.weight.view(self.conv1d.weight.size(0),
273
self.conv1d.weight.size(2))
274
302
- if envs.VLLM_USE_V1 and attn_metadata is None:
303
- # V1 profile run
275
+ if attn_metadata is None:
276
+ # profile run
277
hidden_states = (hidden_states.transpose(0, 1).clone().transpose(
278
0, 1)).contiguous()
279
output[:] = self.out_proj(hidden_states)
@@ -316,42 +289,23 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
289
# NOTE: V0 put prefill before decode, v1 puts decode before prefill
290
# Separate prefill and decode by splitting varlen input
291
# Split along token dimension
319
- if envs.VLLM_USE_V1:
320
- hidden_states_d, hidden_states_p = torch.split(
321
- hidden_states[:num_actual_tokens],
322
- [num_decodes, num_prefill_tokens],
323
- dim=0,
324
- )
325
- gate_d, gate_p = torch.split(gate[:num_actual_tokens],
326
- [num_decodes, num_prefill_tokens],
327
- dim=0)
328
- # Split along batch dimension
329
- state_indices_tensor_d, state_indices_tensor_p = torch.split(
330
- state_indices_tensor,
331
- [num_decodes, num_prefills],
332
- dim=0,
333
- )
334
- query_start_loc_p = (
335
- attn_metadata.query_start_loc[-num_prefills - 1:] -
336
- num_decodes if has_prefill else None)
337
- else:
338
- hidden_states_p, hidden_states_d = torch.split(
339
- hidden_states,
340
- [num_prefill_tokens, num_decodes],
341
- dim=0,
342
- )
343
- gate_p, gate_d = torch.split(gate,
344
- [num_prefill_tokens, num_decodes],
345
- dim=0)
346
- # Split along batch dimension
347
- state_indices_tensor_p, state_indices_tensor_d = torch.split(
348
- state_indices_tensor,
349
- [num_prefills, num_decodes],
350
- dim=0,
351
- )
352
- query_start_loc_p = (attn_metadata.query_start_loc[:num_prefills +
353
- 1]
354
- if has_prefill else None)
292
+ hidden_states_d, hidden_states_p = torch.split(
293
+ hidden_states[:num_actual_tokens],
294
+ [num_decodes, num_prefill_tokens],
295
+ dim=0,
296
+ )
297
+ gate_d, gate_p = torch.split(gate[:num_actual_tokens],
298
+ [num_decodes, num_prefill_tokens],
299
+ dim=0)
300
+ # Split along batch dimension
301
+ state_indices_tensor_d, state_indices_tensor_p = torch.split(
302
+ state_indices_tensor,
303
+ [num_decodes, num_prefills],
304
+ dim=0,
305
+ )
306
+ query_start_loc_p = (
307
+ attn_metadata.query_start_loc[-num_prefills - 1:] -
308
+ num_decodes if has_prefill else None)
309
310
# Preallocate output tensor to avoid memcpy cost for merging prefill
311
# and decode outputs
@@ -363,18 +317,11 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
317
dtype=hidden_states.dtype,
318
device=hidden_states.device,
319
)
366
- if envs.VLLM_USE_V1:
367
- preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(
368
- preallocated_ssm_out,
369
- [num_decodes, num_prefill_tokens],
370
- dim=0,
371
- )
372
- else:
373
- preallocated_ssm_out_p, preallocated_ssm_out_d = torch.split(
374
- preallocated_ssm_out,
375
- [num_prefill_tokens, num_decodes],
376
- dim=0,
377
- )
320
+ preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(
321
+ preallocated_ssm_out,
322
+ [num_decodes, num_prefill_tokens],
323
+ dim=0,
324
+ )
325
326
# Process prefill requests
327
if has_prefill:
@@ -383,9 +330,6 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
330
# pointed to by "state_indices_tensor"
331
x = hidden_states_p.transpose(
332
0, 1) # this is the form that causal-conv see
386
- if mamba2_metadata.cu_seqlen is None:
387
- mamba2_metadata = update_metadata(x, query_start_loc_p,
388
- mamba2_metadata)
333
hidden_states_p = causal_conv1d_fn(
334
x,
335
conv_weights,
@@ -394,7 +338,7 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
338
conv_states=conv_state,
339
has_initial_state=has_initial_states_p,
340
cache_indices=state_indices_tensor_p,
397
- metadata=mamba2_metadata,
341
+ metadata=attn_metadata,
342
query_start_loc=query_start_loc_p)
343
hidden_states_p = hidden_states_p.transpose(0, 1)
344
hidden_states_p = hidden_states_p[:num_prefill_tokens]
@@ -470,7 +414,7 @@ class Plamo2MambaMixer(MambaBase, CustomOp):
414
-1, self.num_heads // self.tp_size, self.head_dim)
415
416
# - the hidden is reshaped into (bs, num_heads, head_dim)
473
- # - mamba_cache_params.ssm_state's slots will be selected
417
+ # - ssm_state's slots will be selected
418
# using state_indices_tensor_d
419
420
# NOTE: final output is an in-place update of out tensor
@@ -530,10 +474,7 @@ def plamo2_mamba_mixer(
474
) -> None:
475
forward_context: ForwardContext = get_forward_context()
476
self = forward_context.no_compile_layers[layer_name]
533
- self.forward_cuda(hidden_states=hidden_states,
534
- output=output,
535
- mamba_cache_params=None,
536
- mamba2_metadata=None)
477
+ self.forward_cuda(hidden_states=hidden_states, output=output)
478
479
480
def plamo2_mamba_mixer_fake(
@@ -731,8 +672,6 @@ class Plamo2DecoderLayer(nn.Module):
672
positions: torch.Tensor,
673
hidden_states: torch.Tensor,
674
residual: Optional[torch.Tensor],
734
- mamba_cache_params: MambaCacheParams,
735
- mamba2_metadata: Mamba2Metadata,
675
**kwargs,
676
):
677
if residual is None:
@@ -747,8 +686,6 @@ class Plamo2DecoderLayer(nn.Module):
686
output = torch.empty_like(hidden_states)
687
mixer_kwargs = {
688
"output": output,
750
- "mamba_cache_params": mamba_cache_params,
751
- "mamba2_metadata": mamba2_metadata,
689
}
690
else:
691
mixer_kwargs = {
@@ -790,23 +727,12 @@ class Plamo2Decoder(torch.nn.Module):
727
positions: torch.Tensor,
728
hidden_states: torch.Tensor,
729
residual: Optional[torch.Tensor],
793
- mamba_cache_params: MambaCacheParams,
794
- mamba2_metadata: Mamba2Metadata,
730
) -> torch.Tensor:
796
- mamba_cache_index = 0
731
for layer in islice(self.layers, self.start_layer, self.end_layer):
798
- layer_mamba_cache_params = None
799
- if layer.is_mamba and mamba_cache_params is not None:
800
- layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
801
- mamba_cache_index)
802
- mamba_cache_index += 1
803
-
732
hidden_states, residual = layer(
733
positions=positions,
734
hidden_states=hidden_states,
735
residual=residual,
808
- mamba_cache_params=layer_mamba_cache_params,
809
- mamba2_metadata=mamba2_metadata,
736
)
737
return hidden_states, residual
738
@@ -844,7 +770,6 @@ class Plamo2Model(torch.nn.Module):
770
self,
771
input_ids: torch.Tensor,
772
positions: torch.Tensor,
847
- mamba_cache_params: MambaCacheParams,
773
intermediate_tensors: Optional[IntermediateTensors] = None,
774
inputs_embeds: Optional[torch.Tensor] = None,
775
) -> torch.Tensor:
@@ -859,23 +784,10 @@ class Plamo2Model(torch.nn.Module):
784
hidden_states = intermediate_tensors["hidden_states"]
785
residual = intermediate_tensors["residual"]
786
862
- if not envs.VLLM_USE_V1:
863
- attn_metadata: AttentionMetadata = get_forward_context(
864
- ).attn_metadata
865
- mamba2_metadata = prepare_mamba2_metadata(
866
- chunk_size=self.config.mamba_chunk_size,
867
- attn_metadata=attn_metadata,
868
- )
869
- else:
870
- # v1 get mamba2_metadata from forward_context
871
- mamba2_metadata = None
872
-
787
hidden_states, residual = self.layers(
788
positions=positions,
789
hidden_states=hidden_states,
790
residual=residual,
877
- mamba_cache_params=mamba_cache_params,
878
- mamba2_metadata=mamba2_metadata,
791
)
792
if not get_pp_group().is_last_rank:
793
return IntermediateTensors({
@@ -925,9 +837,6 @@ class Plamo2ForCausalLM(torch.nn.Module, HasInnerState, SupportsPP, IsHybrid):
837
if self.config.tie_word_embeddings:
838
self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
839
928
- # Used to track and store by the Mamba cache between steps.
929
- self.mamba_cache: Optional[MambaCacheManager] = None
930
-
840
self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
841
self.config.vocab_size)
842
self.make_empty_intermediate_tensors = (
@@ -942,39 +851,11 @@ class Plamo2ForCausalLM(torch.nn.Module, HasInnerState, SupportsPP, IsHybrid):
851
intermediate_tensors: Optional[IntermediateTensors] = None,
852
inputs_embeds: Optional[torch.Tensor] = None,
853
**kwargs):
945
- if not envs.VLLM_USE_V1:
946
- if self.mamba_cache is None:
947
- num_mamba_layers = (
948
- self.model_config.get_num_layers_by_block_type(
949
- self.vllm_config.parallel_config,
950
- LayerBlockType.mamba))
951
-
952
- mamba_state_shape = self.get_mamba_state_shape_from_config(
953
- self.vllm_config, use_v1=False)
954
- mamba_state_dtype = \
955
- self.get_mamba_state_dtype_from_config(
956
- self.vllm_config)
957
- self.mamba_cache = MambaCacheManager(self.vllm_config,
958
- num_mamba_layers,
959
- *mamba_state_shape,
960
- *mamba_state_dtype)
961
-
962
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
963
- else:
964
- # NOTE: mamba_cache_params is not needed for v1
965
- mamba_cache_params = None
854
967
- hidden_states = self.model(input_ids, positions, mamba_cache_params,
968
- intermediate_tensors, inputs_embeds)
855
+ hidden_states = self.model(input_ids, positions, intermediate_tensors,
856
+ inputs_embeds)
857
return hidden_states
858
971
- def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
972
- return self.mamba_cache.copy_inputs_before_cuda_graphs(
973
- input_buffers, **kwargs)
974
-
975
- def get_seqlen_agnostic_capture_inputs(self, batch_size: int):
976
- return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
977
-
859
@classmethod
860
def get_mamba_state_dtype_from_config(
861
cls,
@@ -991,12 +872,10 @@ class Plamo2ForCausalLM(torch.nn.Module, HasInnerState, SupportsPP, IsHybrid):
872
def get_mamba_state_shape_from_config(
873
cls,
874
vllm_config: "VllmConfig",
994
- use_v1: bool = True,
875
) -> tuple[tuple[int, int], tuple[int, int, int]]:
876
"""Calculate shapes for Mamba's convolutional and state caches.
877
Args:
878
vllm_config: vLLM config
999
- use_v1: Get shapes for V1 (or V0)
879
Returns:
880
Tuple containing:
881
- conv_state_shape: Shape for convolutional state cache
@@ -1015,7 +894,6 @@ class Plamo2ForCausalLM(torch.nn.Module, HasInnerState, SupportsPP, IsHybrid):
894
head_dim=hf_config.hidden_size_per_head,
895
state_size=hf_config.mamba_d_state,
896
conv_kernel=hf_config.mamba_d_conv,
1018
- use_v1=use_v1,
897
)
898
899
def compute_logits(
vllm/model_executor/models/qwen3_next.py
+8
-28
@@ -11,7 +11,6 @@ from einops import rearrange
11
from torch import nn
12
from transformers.activations import ACT2FN
13
14
-from vllm import envs
14
from vllm.attention import Attention, AttentionBackend, AttentionMetadata
15
from vllm.compilation.decorators import support_torch_compile
16
from vllm.config import (CacheConfig, ModelConfig, SpeculativeConfig,
@@ -35,7 +34,6 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
34
RowParallelLinear)
35
from vllm.model_executor.layers.logits_processor import LogitsProcessor
36
from vllm.model_executor.layers.mamba.abstract import MambaBase
38
-from vllm.model_executor.layers.mamba.mamba2_metadata import update_metadata
37
from vllm.model_executor.layers.mamba.mamba_mixer2 import (
38
mamba_v2_sharded_weight_loader)
39
from vllm.model_executor.layers.mamba.mamba_utils import (
@@ -51,7 +49,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
49
DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
50
from vllm.model_executor.model_loader.weight_utils import (
51
default_weight_loader, sharded_weight_loader)
54
-from vllm.model_executor.models.mamba_cache import MambaCacheParams
52
from vllm.model_executor.models.qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
53
from vllm.model_executor.utils import set_weight_attrs
54
from vllm.platforms import current_platform
@@ -198,14 +195,8 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
195
196
def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
197
return MambaStateShapeCalculator.gated_delta_net_state_shape(
201
- self.tp_size,
202
- self.num_k_heads,
203
- self.num_v_heads,
204
- self.head_k_dim,
205
- self.head_v_dim,
206
- self.conv_kernel_size,
207
- self.num_spec,
208
- use_v1=True)
198
+ self.tp_size, self.num_k_heads, self.num_v_heads, self.head_k_dim,
199
+ self.head_v_dim, self.conv_kernel_size, self.num_spec)
200
201
def __init__(
202
self,
@@ -394,7 +385,6 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
385
self,
386
hidden_states: torch.Tensor,
387
output: torch.Tensor,
397
- cache_params: Optional[MambaCacheParams] = None,
388
):
389
return torch.ops.vllm.gdn_attention(
390
hidden_states,
@@ -416,7 +406,6 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
406
407
assert isinstance(attn_metadata, dict)
408
attn_metadata = attn_metadata[self.prefix]
419
- conv_metadata = attn_metadata
409
assert isinstance(attn_metadata, GDNAttentionMetadata)
410
has_initial_state = attn_metadata.has_initial_state
411
spec_query_start_loc = attn_metadata.spec_query_start_loc
@@ -479,12 +468,8 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
468
# 2.2: process the remaining part
469
if attn_metadata.num_prefills > 0:
470
mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1)
482
- if conv_metadata.cu_seqlen is None:
483
- conv_metadata = update_metadata(mixed_qkv_non_spec_T,
484
- non_spec_query_start_loc,
485
- conv_metadata)
471
# - "cache_indices" updates the conv_state cache in positions
487
- # pointed to by "mamba_cache_params.state_indices_tensor"
472
+ # pointed to by "state_indices_tensor"
473
mixed_qkv_non_spec = causal_conv1d_fn(
474
mixed_qkv_non_spec_T,
475
conv_weights,
@@ -494,7 +479,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
479
has_initial_state=has_initial_state,
480
cache_indices=non_spec_state_indices_tensor,
481
query_start_loc=non_spec_query_start_loc,
497
- metadata=conv_metadata,
482
+ metadata=attn_metadata,
483
).transpose(0, 1)
484
elif attn_metadata.num_decodes > 0:
485
mixed_qkv_non_spec = causal_conv1d_update(
@@ -1075,7 +1060,6 @@ class Qwen3NextForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
1060
scheduler_config = vllm_config.scheduler_config
1061
assert not cache_config.enable_prefix_caching, \
1062
"Qwen3Next currently does not support prefix caching"
1078
- assert envs.VLLM_USE_V1, "Qwen3Next requires VLLM_USE_V1"
1063
self.quant_config = vllm_config.quant_config
1064
1065
super().__init__()
@@ -1195,14 +1179,10 @@ class Qwen3NextForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP,
1179
num_spec = (vllm_config.speculative_config.num_speculative_tokens
1180
if vllm_config.speculative_config else 0)
1181
return MambaStateShapeCalculator.gated_delta_net_state_shape(
1198
- tp_size,
1199
- hf_config.linear_num_key_heads,
1200
- hf_config.linear_num_value_heads,
1201
- hf_config.linear_key_head_dim,
1202
- hf_config.linear_value_head_dim,
1203
- hf_config.linear_conv_kernel_dim,
1204
- num_spec,
1205
- use_v1=True)
1182
+ tp_size, hf_config.linear_num_key_heads,
1183
+ hf_config.linear_num_value_heads, hf_config.linear_key_head_dim,
1184
+ hf_config.linear_value_head_dim, hf_config.linear_conv_kernel_dim,
1185
+ num_spec)
1186
1187
def compute_logits(
1188
self,
vllm/model_executor/models/registry.py
-1
@@ -134,7 +134,6 @@ _TEXT_GENERATION_MODELS = {
134
"PhiForCausalLM": ("phi", "PhiForCausalLM"),
135
"Phi3ForCausalLM": ("phi3", "Phi3ForCausalLM"),
136
"PhiMoEForCausalLM": ("phimoe", "PhiMoEForCausalLM"),
137
- "Phi4FlashForCausalLM": ("phi4flash", "Phi4FlashForCausalLM"),
137
"Plamo2ForCausalLM": ("plamo2", "Plamo2ForCausalLM"),
138
"QWenLMHeadModel": ("qwen", "QWenLMHeadModel"),
139
"Qwen2ForCausalLM": ("qwen2", "Qwen2ForCausalLM"),
vllm/model_executor/models/zamba2.py
-93
@@ -15,12 +15,10 @@ import torch
15
from torch import nn
16
from transformers import Zamba2Config
17
18
-from vllm import envs
18
from vllm.attention.layer import Attention
19
from vllm.compilation.decorators import support_torch_compile
20
from vllm.config import CacheConfig, ModelConfig, VllmConfig
21
from vllm.distributed import get_tensor_model_parallel_world_size
23
-from vllm.forward_context import get_forward_context
22
from vllm.model_executor.layers.activation import GeluAndMul
23
from vllm.model_executor.layers.layernorm import RMSNorm
24
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
@@ -29,8 +27,6 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
27
ReplicatedLinear,
28
RowParallelLinear)
29
from vllm.model_executor.layers.logits_processor import LogitsProcessor
32
-from vllm.model_executor.layers.mamba.mamba2_metadata import (
33
- Mamba2Metadata, prepare_mamba2_metadata)
30
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
31
from vllm.model_executor.layers.mamba.mamba_utils import (
32
MambaStateDtypeCalculator, MambaStateShapeCalculator)
@@ -39,8 +35,6 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
35
from vllm.model_executor.layers.vocab_parallel_embedding import (
36
DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)
37
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
42
-from vllm.model_executor.models.mamba_cache import (MambaCacheManager,
43
- MambaCacheParams)
38
from vllm.sequence import IntermediateTensors
39
40
from .interfaces import HasInnerState, IsHybrid
@@ -515,8 +509,6 @@ class Zamba2MambaDecoderLayer(nn.Module):
509
def forward(
510
self,
511
hidden_states: torch.Tensor,
518
- mamba_cache_params: MambaCacheParams,
519
- 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,
@@ -525,8 +517,6 @@ class Zamba2MambaDecoderLayer(nn.Module):
517
518
Args:
519
hidden_states: Input tensor [batch_size, seq_len, hidden_size]
528
- mamba_cache_params: Parameters for Mamba's state caches
529
- (one for conv, one for ssm)
520
transformer_hidden_states: Optional output from transformer path
521
Added to input if provided (used in hybrid architecture)
522
positions: Optional position IDs (unused in Mamba)
@@ -555,8 +545,6 @@ class Zamba2MambaDecoderLayer(nn.Module):
545
self.mamba(
546
hidden_states,
547
output,
558
- mamba_cache_params=mamba_cache_params,
559
- mamba2_metadata=mamba2_metadata,
548
)
549
550
# residual connection after mamba
@@ -607,8 +595,6 @@ class Zamba2HybridLayer(nn.Module):
595
hidden_states: torch.Tensor,
596
original_hidden_states: torch.Tensor,
597
positions: torch.Tensor,
610
- mamba_cache_params: MambaCacheParams,
611
- mamba2_metadata: Mamba2Metadata,
598
) -> torch.Tensor:
599
"""Forward pass through the hybrid layer.
600
@@ -623,8 +609,6 @@ class Zamba2HybridLayer(nn.Module):
609
original_hidden_states: Original input for transformer residual
610
connection
611
positions: Position IDs for positional embeddings
626
- mamba_cache_params: Parameters for Mamba's state caches
627
- (one for conv, one for ssm)
612
613
Returns:
614
Output tensor combining transformer and Mamba representations
@@ -644,8 +628,6 @@ class Zamba2HybridLayer(nn.Module):
628
layer_outputs = self.mamba_decoder(
629
hidden_states,
630
transformer_hidden_states=transformer_hidden_states,
647
- mamba_cache_params=mamba_cache_params,
648
- mamba2_metadata=mamba2_metadata,
631
)
632
633
return layer_outputs
@@ -752,7 +734,6 @@ class Zamba2Model(nn.Module):
734
self,
735
input_ids: torch.Tensor,
736
positions: torch.Tensor,
755
- mamba_cache_params: MambaCacheParams,
737
inputs_embeds: Optional[torch.Tensor] = None,
738
) -> Union[torch.Tensor, IntermediateTensors]:
739
"""Forward pass through the model.
@@ -760,8 +741,6 @@ class Zamba2Model(nn.Module):
741
Args:
742
input_ids: Input token IDs
743
positions: Position IDs for embeddings
763
- mamba_cache_params: Parameters for Mamba's state caches
764
- (one for conv, one for ssm)
744
inputs_embeds: Optional pre-computed input embeddings
745
746
Returns:
@@ -773,33 +752,13 @@ class Zamba2Model(nn.Module):
752
inputs_embeds = self.get_input_embeddings(input_ids)
753
hidden_states = inputs_embeds
754
776
- attn_metadata = get_forward_context().attn_metadata
777
-
778
- if not envs.VLLM_USE_V1:
779
- mamba2_metadata = prepare_mamba2_metadata(
780
- chunk_size=self.config.chunk_size,
781
- attn_metadata=attn_metadata,
782
- )
783
- else:
784
- # v1 get mamba2_metadata from forward_context
785
- mamba2_metadata = None
786
-
755
# Process through layers
756
original_hidden_states = torch.clone(hidden_states)
757
for layer_idx, layer in enumerate(self.layers):
790
-
791
- layer_mamba_cache_params = None
792
- if (isinstance(layer, (Zamba2HybridLayer, Zamba2MambaDecoderLayer))
793
- and mamba_cache_params):
794
- layer_mamba_cache_params = mamba_cache_params.at_layer_idx(
795
- layer_idx)
796
-
758
layer_outputs = layer(
759
hidden_states,
760
original_hidden_states=original_hidden_states,
761
positions=positions,
801
- mamba_cache_params=layer_mamba_cache_params,
802
- mamba2_metadata=mamba2_metadata,
762
)
763
hidden_states = layer_outputs
764
@@ -870,13 +829,11 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
829
def get_mamba_state_shape_from_config(
830
cls,
831
vllm_config: "VllmConfig",
873
- use_v1: bool = True,
832
) -> tuple[tuple[int, int], tuple[int, int, int]]:
833
"""Calculate shapes for Mamba's convolutional and state caches.
834
835
Args:
836
vllm_config: vLLM config
879
- use_v1: Get shapes for V1 (or V0)
837
838
Returns:
839
Tuple containing:
@@ -896,7 +853,6 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
853
head_dim=hf_config.mamba_headdim,
854
state_size=hf_config.mamba_d_state,
855
conv_kernel=hf_config.mamba_d_conv,
899
- use_v1=use_v1,
856
)
857
858
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
@@ -945,9 +901,6 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
901
# Tie weights with input embeddings if using same dimensions
902
self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
903
948
- # Used to track and store by the Mamba cache between steps.
949
- self.mamba_cache: Optional[MambaCacheManager] = None
950
-
904
# Initialize logits processing and sampling
905
self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,
906
config.vocab_size)
@@ -977,61 +930,15 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
930
Returns:
931
Output hidden states
932
"""
980
- # Initialize Mamba cache if needed
981
- mamba_cache_params = None
982
- if not envs.VLLM_USE_V1:
983
- if self.mamba_cache is None:
984
- num_mamba_layers = self.config.num_hidden_layers
985
- mamba_state_shape = \
986
- self.get_mamba_state_shape_from_config(
987
- self.vllm_config, use_v1=False)
988
- mamba_state_dtype = \
989
- self.get_mamba_state_dtype_from_config(
990
- self.vllm_config)
991
- self.mamba_cache = MambaCacheManager(self.vllm_config,
992
- num_mamba_layers,
993
- *mamba_state_shape,
994
- *mamba_state_dtype)
995
-
996
- # Get cache parameters for current run
997
- mamba_cache_params = self.mamba_cache.current_run_tensors(**kwargs)
998
-
933
# Forward pass through model
934
hidden_states = self.model(
935
input_ids,
936
positions,
1003
- mamba_cache_params,
937
inputs_embeds,
938
)
939
940
return hidden_states
941
1009
- def copy_inputs_before_cuda_graphs(
1010
- self, input_buffers: dict[str, torch.Tensor],
1011
- **kwargs: Any) -> dict[str, torch.Tensor]:
1012
- """Copy inputs before CUDA graph capture.
1013
-
1014
- Args:
1015
- input_buffers: Dictionary of input tensors
1016
- **kwargs: Additional arguments passed to cache manager
1017
-
1018
- Returns:
1019
- Updated input buffers
1020
- """
1021
- return self.mamba_cache.copy_inputs_before_cuda_graphs(
1022
- input_buffers, **kwargs)
1023
-
1024
- def get_seqlen_agnostic_capture_inputs(
1025
- self, batch_size: int) -> dict[str, torch.Tensor]:
1026
- """Get inputs for sequence-length-agnostic graph capture.
1027
-
1028
- Args:
1029
- batch_size: Size of batch to capture
1030
- Returns:
1031
- Dictionary of capture inputs
1032
- """
1033
- return self.mamba_cache.get_seqlen_agnostic_capture_inputs(batch_size)
1034
-
942
def compute_logits(
943
self,
944
hidden_states: torch.Tensor,
vllm/v1/attention/backends/gdn_attn.py
+7
-1
@@ -12,6 +12,7 @@ from vllm.config import VllmConfig
12
from vllm.v1.attention.backends.utils import (AttentionCGSupport,
13
AttentionMetadataBuilder,
14
CommonAttentionMetadata,
15
+ compute_causal_conv1d_metadata,
16
split_decodes_and_prefills)
17
from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
18
@@ -52,7 +53,6 @@ class GDNAttentionMetadata:
53
54
# The following attributes are for triton implementation of causal_conv1d
55
nums_dict: Optional[dict] = None
55
- cu_seqlen: Optional[int] = None
56
batch_ptr: Optional[torch.Tensor] = None
57
token_chunk_offset_ptr: Optional[torch.Tensor] = None
58
@@ -134,6 +134,7 @@ class GDNAttentionMetadataBuilder(
134
context_lens = m.num_computed_tokens_cpu
135
context_lens_tensor = context_lens.to(query_start_loc.device)
136
seq_lens_tensor = m.seq_lens
137
+ nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
138
139
if (not self.use_spec_decode or num_draft_tokens is None
140
or num_draft_tokens.sum().item() == 0):
@@ -210,6 +211,8 @@ class GDNAttentionMetadataBuilder(
211
has_initial_state = context_lens_tensor > 0
212
if spec_sequence_masks is not None:
213
has_initial_state = has_initial_state[~spec_sequence_masks]
214
+ nums_dict, batch_ptr, token_chunk_offset_ptr = \
215
+ compute_causal_conv1d_metadata(non_spec_query_start_loc)
216
else:
217
has_initial_state = None
218
num_actual_tokens = num_prefill_tokens + num_decode_tokens + \
@@ -297,6 +300,9 @@ class GDNAttentionMetadataBuilder(
300
spec_sequence_masks=spec_sequence_masks,
301
spec_token_masks=spec_token_masks,
302
num_accepted_tokens=num_accepted_tokens,
303
+ nums_dict=nums_dict,
304
+ batch_ptr=batch_ptr,
305
+ token_chunk_offset_ptr=token_chunk_offset_ptr,
306
)
307
return attn_metadata
308
vllm/v1/attention/backends/mamba2_attn.py
+12
-3
@@ -7,11 +7,12 @@ from typing import Optional
7
import torch
8
9
from vllm.attention.backends.abstract import AttentionBackend
10
-from vllm.attention.backends.utils import PAD_SLOT_ID
10
from vllm.config import VllmConfig
11
from vllm.v1.attention.backends.mamba_attn import (
12
BaseMambaAttentionMetadataBuilder)
14
-from vllm.v1.attention.backends.utils import (CommonAttentionMetadata,
13
+from vllm.v1.attention.backends.utils import (PAD_SLOT_ID,
14
+ CommonAttentionMetadata,
15
+ compute_causal_conv1d_metadata,
16
split_decodes_and_prefills)
17
from vllm.v1.kv_cache_interface import AttentionSpec
18
@@ -131,7 +132,6 @@ class Mamba2AttentionMetadata:
132
133
# The following attributes are for triton implementation of causal_conv1d
134
nums_dict: Optional[dict] = None
134
- cu_seqlen: Optional[int] = None
135
batch_ptr: Optional[torch.Tensor] = None
136
token_chunk_offset_ptr: Optional[torch.Tensor] = None
137
@@ -161,6 +161,9 @@ class Mamba2AttentionMetadataBuilder(
161
has_initial_states_p = None
162
prep_initial_states = False
163
164
+ # for causal_conv1d
165
+ nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
166
+
167
state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
168
169
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
@@ -198,6 +201,9 @@ class Mamba2AttentionMetadataBuilder(
201
query_start_loc_p, self.chunk_size,
202
num_prefill_tokens))
203
204
+ nums_dict, batch_ptr, token_chunk_offset_ptr = \
205
+ compute_causal_conv1d_metadata(query_start_loc_p)
206
+
207
elif num_decodes <= self.decode_cudagraph_max_bs:
208
# Pad state tensor for CUDA graph
209
num_input_tokens = self.vllm_config.pad_for_cudagraph(num_decodes)
@@ -220,5 +226,8 @@ class Mamba2AttentionMetadataBuilder(
226
chunk_indices_p=chunk_indices_p,
227
chunk_offsets_p=chunk_offsets_p,
228
state_indices_tensor=state_indices_tensor,
229
+ nums_dict=nums_dict,
230
+ batch_ptr=batch_ptr,
231
+ token_chunk_offset_ptr=token_chunk_offset_ptr,
232
)
233
return attn_metadata
vllm/v1/attention/backends/short_conv_attn.py
+13
-1
@@ -9,6 +9,7 @@ from vllm.attention.backends.abstract import AttentionBackend
9
from vllm.config import VllmConfig
10
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
11
CommonAttentionMetadata,
12
+ compute_causal_conv1d_metadata,
13
split_decodes_and_prefills)
14
from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
15
@@ -33,7 +34,6 @@ class ShortConvAttentionMetadata:
34
35
# For causal_conv1d
36
nums_dict: Optional[dict] = None
36
- cu_seqlen: Optional[int] = None
37
batch_ptr: Optional[torch.Tensor] = None
38
token_chunk_offset_ptr: Optional[torch.Tensor] = None
39
@@ -57,6 +57,9 @@ class ShortConvAttentionMetadataBuilder(
57
58
state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
59
60
+ # for causal_conv1d
61
+ nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
62
+
63
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
64
split_decodes_and_prefills(
65
common_attn_metadata,
@@ -70,6 +73,12 @@ class ShortConvAttentionMetadataBuilder(
73
has_initial_states = has_initial_states_cpu.to(
74
query_start_loc.device)
75
76
+ query_start_loc_p = common_attn_metadata.query_start_loc[
77
+ -num_prefills - 1:] - num_decode_tokens
78
+
79
+ nums_dict, batch_ptr, token_chunk_offset_ptr = \
80
+ compute_causal_conv1d_metadata(query_start_loc_p)
81
+
82
attn_metadata = ShortConvAttentionMetadata(
83
num_prefills=num_prefills,
84
num_prefill_tokens=num_prefill_tokens,
@@ -78,5 +87,8 @@ class ShortConvAttentionMetadataBuilder(
87
query_start_loc=query_start_loc,
88
has_initial_states=has_initial_states,
89
state_indices_tensor=state_indices_tensor,
90
+ nums_dict=nums_dict,
91
+ batch_ptr=batch_ptr,
92
+ token_chunk_offset_ptr=token_chunk_offset_ptr,
93
)
94
return attn_metadata
vllm/v1/attention/backends/utils.py
+51
@@ -34,6 +34,8 @@ logger = init_logger(__name__)
34
KVCacheLayoutType = Literal["NHD", "HND"]
35
_KV_CACHE_LAYOUT_OVERRIDE: Union[KVCacheLayoutType, None] = None
36
37
+PAD_SLOT_ID = -1
38
+
39
40
def is_valid_kv_cache_layout(value: str) -> bool:
41
return value in get_args(KVCacheLayoutType)
@@ -838,3 +840,52 @@ def create_fast_prefill_custom_backend(
840
builder_cls=FastPrefillAttentionBuilder)
841
842
return attn_backend
843
+
844
+
845
+def compute_causal_conv1d_metadata(query_start_loc_p: torch.Tensor):
846
+
847
+ # Needed for causal_conv1d
848
+ seqlens = query_start_loc_p.diff().to('cpu')
849
+ nums_dict = {} # type: ignore
850
+ batch_ptr = None
851
+ token_chunk_offset_ptr = None
852
+ for BLOCK_M in [8]: # cover all BLOCK_M values
853
+ nums = -(-seqlens // BLOCK_M)
854
+ nums_dict[BLOCK_M] = {}
855
+ nums_dict[BLOCK_M]['nums'] = nums
856
+ nums_dict[BLOCK_M]['tot'] = nums.sum().item()
857
+ mlist = torch.from_numpy(np.repeat(np.arange(len(nums)), nums))
858
+ nums_dict[BLOCK_M]['mlist'] = mlist
859
+ mlist_len = len(nums_dict[BLOCK_M]['mlist'])
860
+ nums_dict[BLOCK_M]['mlist_len'] = mlist_len
861
+ MAX_NUM_PROGRAMS = max(1024, mlist_len) * 2
862
+ offsetlist = [] # type: ignore
863
+ for idx, num in enumerate(nums):
864
+ offsetlist.extend(range(num))
865
+ offsetlist = torch.tensor(offsetlist, dtype=torch.int32)
866
+ nums_dict[BLOCK_M]['offsetlist'] = offsetlist
867
+
868
+ if batch_ptr is None:
869
+ # Update default value after class definition
870
+ batch_ptr = torch.full((MAX_NUM_PROGRAMS, ),
871
+ PAD_SLOT_ID,
872
+ dtype=torch.int32,
873
+ device='cuda')
874
+ token_chunk_offset_ptr = torch.full((MAX_NUM_PROGRAMS, ),
875
+ PAD_SLOT_ID,
876
+ dtype=torch.int32,
877
+ device='cuda')
878
+ else:
879
+ if batch_ptr.nelement() < MAX_NUM_PROGRAMS:
880
+ batch_ptr.resize_(MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID)
881
+ token_chunk_offset_ptr.resize_( # type: ignore
882
+ MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID)
883
+
884
+ batch_ptr[0:mlist_len].copy_(mlist)
885
+ token_chunk_offset_ptr[ # type: ignore
886
+ 0:mlist_len].copy_(offsetlist)
887
+ nums_dict[BLOCK_M]['batch_ptr'] = batch_ptr
888
+ nums_dict[BLOCK_M]['token_chunk_offset_ptr'] = (token_chunk_offset_ptr
889
+ ) # type: ignore
890
+
891
+ return nums_dict, batch_ptr, token_chunk_offset_ptr