@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 147 lines 4.15 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from __future__ import annotations
5
6 import torch
7 import torch.nn.functional as F
8
9 from vllm._custom_ops import causal_conv1d_update_cpu_vec
10 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID
11
12
13 def causal_conv1d_fn_cpu(
14 x: torch.Tensor,
15 weight: torch.Tensor,
16 bias: torch.Tensor | None,
17 conv_states: torch.Tensor,
18 query_start_loc: torch.Tensor,
19 cache_indices: torch.Tensor | None = None,
20 has_initial_state: torch.Tensor | None = None,
21 activation: str | None = "silu",
22 pad_slot_id: int = PAD_SLOT_ID,
23 **kwargs,
24 ) -> torch.Tensor:
25 """CPU implementation for causal_conv1d_fwd."""
26 if isinstance(activation, bool) and activation:
27 activation = "silu"
28 elif isinstance(activation, bool):
29 activation = None
30
31 original_x_dtype = x.dtype
32 x = x.to(conv_states.dtype)
33
34 out = torch.empty_like(x)
35 state_len = weight.shape[1] - 1
36 assert activation in {None, "silu", "swish"}
37
38 seq_begin_end_idx = [
39 (int(query_start_loc[idx].item()), int(query_start_loc[idx + 1].item()))
40 for idx in range(query_start_loc.shape[0] - 1)
41 ]
42 weight = weight.unsqueeze(1)
43
44 for seq_idx, (bos, eos) in enumerate(seq_begin_end_idx):
45 if bos == eos:
46 continue
47
48 slot = (
49 int(cache_indices[seq_idx].item()) if cache_indices is not None else seq_idx
50 )
51
52 if slot == pad_slot_id:
53 continue
54
55 seq_x = x[:, bos:eos].unsqueeze(0)
56
57 if has_initial_state is not None and bool(has_initial_state[seq_idx].item()):
58 initial_state = conv_states[slot, :, :state_len].unsqueeze(0)
59 else:
60 initial_state = torch.zeros(
61 1,
62 weight.shape[0],
63 state_len,
64 device=seq_x.device,
65 dtype=seq_x.dtype,
66 )
67
68 conv_input = torch.cat([initial_state, seq_x], dim=-1).to(weight.dtype)
69 seq_out = F.conv1d(
70 conv_input,
71 weight,
72 bias,
73 padding=0,
74 groups=weight.shape[0],
75 )
76 seq_out = seq_out[..., -seq_x.shape[-1] :].to(dtype=x.dtype)
77
78 if activation in ("silu", "swish"):
79 seq_out = F.silu(seq_out)
80
81 out[:, bos:eos] = seq_out.squeeze(0)
82 conv_states[slot, :, :state_len].copy_(conv_input[..., -state_len:].squeeze(0))
83
84 return out.to(original_x_dtype)
85
86
87 def causal_conv1d_update_cpu(
88 x: torch.Tensor,
89 conv_state: torch.Tensor,
90 weight: torch.Tensor,
91 bias: torch.Tensor | None = None,
92 activation: bool | str | None = None,
93 conv_state_indices: torch.Tensor | None = None,
94 query_start_loc: torch.Tensor | None = None,
95 pad_slot_id: int | None = None,
96 **kwargs,
97 ) -> torch.Tensor:
98 """CPU implementation for causal_conv1d_update."""
99 if isinstance(activation, bool):
100 activation = "silu" if activation else None
101
102 if pad_slot_id is None:
103 pad_slot_id = kwargs.get("null_block_id", NULL_BLOCK_ID)
104 if pad_slot_id is None:
105 pad_slot_id = NULL_BLOCK_ID
106
107 return causal_conv1d_update_cpu_vec(
108 x,
109 conv_state,
110 weight,
111 bias,
112 activation,
113 conv_state_indices,
114 query_start_loc,
115 pad_slot_id,
116 )
117
118
119 def causal_conv1d_update_torch(
120 x: torch.Tensor,
121 conv_state: torch.Tensor,
122 weight: torch.Tensor,
123 bias: torch.Tensor | None = None,
124 activation: str | None = None,
125 ) -> torch.Tensor:
126 """Pure PyTorch fallback for causal_conv1d_update.
127 Currently used as a fallback for Arm (aarch64) to leverage
128 oneDNN/ACL F.conv1d kernels for batched decoding.
129 """
130 assert activation in {None, "silu", "swish"}
131
132 _, dim, seq_len = x.shape
133 state_len = conv_state.shape[-1]
134
135 x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype)
136 conv_state.copy_(x_new[:, :, -state_len:])
137
138 out = F.conv1d(
139 x_new,
140 weight.unsqueeze(1),
141 bias,
142 padding=0,
143 groups=dim,
144 )[:, :, -seq_len:]
145 if activation in ("silu", "swish"):
146 out = F.silu(out)
147 return out