|
7 | 7 | # and https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/ops/causal_conv1d.py |
8 | 8 | # mypy: ignore-errors |
9 | 9 |
|
10 | | -from typing import Any |
11 | | - |
12 | 10 | import torch |
13 | | -import torch.nn.functional as F |
14 | | -from vllm.distributed import get_pcp_group |
15 | | -from vllm.forward_context import get_forward_context |
16 | 11 | from vllm.triton_utils import HAS_TRITON, tl, triton |
17 | 12 | from vllm.v1.attention.backends.utils import PAD_SLOT_ID # type: ignore |
18 | 13 |
|
|
26 | 21 | _pytorch_update = None |
27 | 22 |
|
28 | 23 |
|
29 | | -def causal_conv1d_ref( |
30 | | - x: torch.Tensor, |
31 | | - weight: torch.Tensor, |
32 | | - bias: torch.Tensor | None = None, |
33 | | - initial_states: torch.Tensor | None = None, |
34 | | - return_final_states: bool = False, |
35 | | - final_states_out: torch.Tensor | None = None, |
36 | | - activation: str | None = "silu", |
37 | | -): |
38 | | - """ |
39 | | - x: (batch, dim, seqlen) |
40 | | - weight: (dim, width) |
41 | | - bias: (dim,) |
42 | | - initial_states: (batch, dim, width - 1) |
43 | | - final_states_out: (batch, dim, width - 1) |
44 | | - out: (batch, dim, seqlen) |
45 | | - """ |
46 | | - if activation not in [None, "silu", "swish"]: |
47 | | - raise NotImplementedError(f"causal_conv1d_ref activation must be None, silu, or swish, got {activation}") |
48 | | - dtype_in = x.dtype |
49 | | - x = x.to(weight.dtype) |
50 | | - seqlen = x.shape[-1] |
51 | | - dim, width = weight.shape |
52 | | - |
53 | | - if initial_states is None: |
54 | | - out = F.conv1d(x, weight.unsqueeze(1), bias, padding=width - 1, groups=dim) |
55 | | - else: |
56 | | - x = torch.cat([initial_states, x], dim=-1) |
57 | | - out = F.conv1d(x, weight.unsqueeze(1), bias, padding=0, groups=dim) |
58 | | - out = out[..., :seqlen] |
59 | | - |
60 | | - if return_final_states: |
61 | | - final_states = F.pad(x, (width - 1 - x.shape[-1], 0)).to(dtype_in) # (batch, dim, width - 1) |
62 | | - if final_states_out is not None: |
63 | | - final_states_out.copy_(final_states) |
64 | | - else: |
65 | | - final_states_out = final_states |
66 | | - out = (out if activation is None else F.silu(out)).to(dtype=dtype_in) |
67 | | - return (out, None) if not return_final_states else (out, final_states_out) |
68 | | - |
69 | | - |
70 | | -def causal_conv1d_fn( |
71 | | - x: torch.Tensor, |
72 | | - weight: torch.Tensor, |
73 | | - bias: torch.Tensor | None = None, |
74 | | - activation: str | None = "silu", |
75 | | - conv_states: torch.Tensor | None = None, |
76 | | - has_initial_state: torch.Tensor | None = None, |
77 | | - cache_indices: torch.Tensor | None = None, |
78 | | - query_start_loc: torch.Tensor | None = None, |
79 | | - metadata: Any | None = None, |
80 | | - pad_slot_id: int = PAD_SLOT_ID, |
81 | | -): |
82 | | - """ |
83 | | - x: (batch, dim, seqlen) or (dim,cu_seq_len) for varlen |
84 | | - sequences are concatenated from left to right for varlen |
85 | | - weight: (dim, width) |
86 | | - bias: (dim,) |
87 | | - query_start_loc: (batch + 1) int32 |
88 | | - The cumulative sequence lengths of the sequences in |
89 | | - the batch, used to index into sequence. prepended by 0. |
90 | | - for example: query_start_loc = torch.Tensor([0,10,16,17]), |
91 | | - x.shape=(dim,17) |
92 | | - cache_indices: (batch) int32 |
93 | | - indicates the corresponding state index, |
94 | | - like so: conv_state = conv_states[cache_indices[batch_id]] |
95 | | - has_initial_state: (batch) bool |
96 | | - indicates whether should the kernel take the current state as initial |
97 | | - state for the calculations |
98 | | - conv_states: (...,dim,width - 1) itype |
99 | | - updated inplace if provided |
100 | | - activation: either None or "silu" or "swish" |
101 | | - pad_slot_id: int |
102 | | - if cache_indices is passed, lets the kernel identify padded |
103 | | - entries that will not be processed, |
104 | | - for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] |
105 | | - in this case, the kernel will not process entries at |
106 | | - indices 0 and 3 |
107 | | - out: (batch, dim, seqlen) |
108 | | - """ |
109 | | - forward_context = get_forward_context() |
110 | | - num_prefills = 0 |
111 | | - attn_metadata = forward_context.attn_metadata |
112 | | - if attn_metadata is not None and isinstance(attn_metadata, dict): |
113 | | - attn_metadata = next(iter(attn_metadata.values()), None) |
114 | | - if attn_metadata is not None: |
115 | | - num_prefills = attn_metadata.num_prefills |
116 | | - |
117 | | - if activation not in [None, "silu", "swish"]: |
118 | | - raise NotImplementedError(f"causal_conv1d_fn: activation must be None, silu, or swish, got {activation}") |
119 | | - if x.stride(-1) != 1: |
120 | | - x = x.contiguous() |
121 | | - bias = bias.contiguous() if bias is not None else None |
122 | | - |
123 | | - out_ref = [] |
124 | | - out_ref_b = [] |
125 | | - seqlens = query_start_loc[1:] - query_start_loc[:-1] |
126 | | - seqlens = seqlens.tolist() |
127 | | - splits = torch.split(x, seqlens, dim=-1) |
128 | | - width = weight.shape[1] |
129 | | - state_len = width - 1 |
130 | | - prefill_seq_offset = max(0, len(seqlens) - num_prefills) |
131 | | - last_width_prefill_x = extract_last_width(x, query_start_loc[prefill_seq_offset:], state_len) |
132 | | - |
133 | | - pcp_rank = get_pcp_group().rank_in_group |
134 | | - if get_pcp_group().world_size > 1: |
135 | | - all_last_width_prefill_x = get_pcp_group().all_gather(last_width_prefill_x.unsqueeze(0).contiguous(), 0) |
136 | | - if pcp_rank > 0: |
137 | | - conv_states[cache_indices[prefill_seq_offset:], :, :state_len] = all_last_width_prefill_x[pcp_rank - 1, ...] |
138 | | - |
139 | | - for i in range(len(seqlens)): |
140 | | - x_s = splits[i] |
141 | | - if cache_indices[i] == PAD_SLOT_ID: |
142 | | - continue |
143 | | - if pcp_rank == 0 and not bool(has_initial_state[i].item()): |
144 | | - initial_states = None |
145 | | - else: |
146 | | - initial_states = conv_states[cache_indices[i]][..., : (width - 1)] |
147 | | - out_ref_b.append( |
148 | | - causal_conv1d_ref( |
149 | | - x_s, |
150 | | - weight, |
151 | | - bias, |
152 | | - activation=activation, |
153 | | - return_final_states=True, |
154 | | - final_states_out=conv_states[cache_indices[i]][..., : (width - 1)].unsqueeze(0), |
155 | | - initial_states=initial_states, |
156 | | - ) |
157 | | - ) |
158 | | - |
159 | | - if get_pcp_group().world_size > 1: |
160 | | - conv_states[cache_indices[prefill_seq_offset:], :, :state_len] = all_last_width_prefill_x[-1, ...] |
161 | | - out_ref.append(torch.cat([t[0] for t in out_ref_b], dim=-1)) |
162 | | - out_ref_tensor = torch.cat(out_ref, dim=0) |
163 | | - return out_ref_tensor |
164 | | - |
165 | | - |
166 | 24 | def extract_last_width(x, start_loc, width): |
167 | 25 | end_loc = start_loc[1:] |
168 | 26 | offsets = torch.arange(width, device=x.device) |
|
0 commit comments