kvcacheio.py 5.86 KB
Newer Older
1
2
from typing import List

3
4
5
import torch


6
7
8
9
10
11
12
def is_hip() -> bool:
    return torch.version.hip is not None


_is_hip = is_hip()


liucong's avatar
liucong committed
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def dcu_alloc_decode_kernel(
    seq_lens_ptr: torch.Tensor,   
    last_loc_ptr: torch.Tensor,    
    free_page_ptr: torch.Tensor ,   
    out_indices: torch.Tensor , 
    bs: int,          
    bs_upper: int,              
    page_size: int,              
):
    torch.ops.sgl_kernel.dcu_alloc_decode_kernel(
        seq_lens_ptr,
        last_loc_ptr,
        free_page_ptr,
        out_indices,
        bs,
        bs_upper,
        page_size,
    )

32
33
34
35
36
37
38
39
40
def transfer_kv_per_layer(
    src_k: torch.Tensor,
    dst_k: torch.Tensor,
    src_v: torch.Tensor,
    dst_v: torch.Tensor,
    src_indices: torch.Tensor,
    dst_indices: torch.Tensor,
    item_size: int,
    block_quota: int = 2,
41
    num_warps_per_block: int = 16 if _is_hip else 32,
42
):
Zhiqiang Xie's avatar
Zhiqiang Xie committed
43
44
45
46
47
48
49
50
51
52
53
    torch.ops.sgl_kernel.transfer_kv_per_layer(
        src_k,
        dst_k,
        src_v,
        dst_v,
        src_indices,
        dst_indices,
        item_size,
        block_quota,
        num_warps_per_block,
    )
54
55


56
def transfer_kv_per_layer_pf_lf(
57
58
59
60
61
62
    src_k: torch.Tensor,
    dst_k: torch.Tensor,
    src_v: torch.Tensor,
    dst_v: torch.Tensor,
    src_indices: torch.Tensor,
    dst_indices: torch.Tensor,
63
    layer_id: int,
64
65
66
    item_size: int,
    src_layout_dim: int,
    block_quota: int = 2,
67
    num_warps_per_block: int = 16 if _is_hip else 32,
68
69
70
71
72
73
74
75
):
    torch.ops.sgl_kernel.transfer_kv_per_layer_pf_lf(
        src_k,
        dst_k,
        src_v,
        dst_v,
        src_indices,
        dst_indices,
76
        layer_id,
77
78
79
80
81
82
83
84
85
86
87
88
89
90
        item_size,
        src_layout_dim,
        block_quota,
        num_warps_per_block,
    )


def transfer_kv_all_layer(
    src_k_layers: torch.Tensor,
    dst_k_layers: torch.Tensor,
    src_v_layers: torch.Tensor,
    dst_v_layers: torch.Tensor,
    src_indices: torch.Tensor,
    dst_indices: torch.Tensor,
91
92
93
    item_size: int,
    num_layers: int,
    block_quota: int = 2,
94
    num_warps_per_block: int = 16 if _is_hip else 32,
95
):
Zhiqiang Xie's avatar
Zhiqiang Xie committed
96
97
98
99
100
101
102
103
104
105
106
107
    torch.ops.sgl_kernel.transfer_kv_all_layer(
        src_k_layers,
        dst_k_layers,
        src_v_layers,
        dst_v_layers,
        src_indices,
        dst_indices,
        item_size,
        num_layers,
        block_quota,
        num_warps_per_block,
    )
108
109


110
111
112
113
114
115
116
117
118
119
120
def transfer_kv_all_layer_lf_pf(
    src_k_layers: torch.Tensor,
    dst_k: torch.Tensor,
    src_v_layers: torch.Tensor,
    dst_v: torch.Tensor,
    src_indices: torch.Tensor,
    dst_indices: torch.Tensor,
    item_size: int,
    dst_layout_dim: int,
    num_layers: int,
    block_quota: int = 2,
121
    num_warps_per_block: int = 16 if _is_hip else 32,
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
):
    torch.ops.sgl_kernel.transfer_kv_all_layer_lf_pf(
        src_k_layers,
        dst_k,
        src_v_layers,
        dst_v,
        src_indices,
        dst_indices,
        item_size,
        dst_layout_dim,
        num_layers,
        block_quota,
        num_warps_per_block,
    )


def transfer_kv_direct(
    src_layers: List[torch.Tensor],
    dst_layers: List[torch.Tensor],
    src_indices: torch.Tensor,
    dst_indices: torch.Tensor,
    page_size: int,
):
    torch.ops.sgl_kernel.transfer_kv_direct(
        src_layers, dst_layers, src_indices, dst_indices, page_size
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
    )


def transfer_kv_per_layer_direct_pf_lf(
    src_ptrs: List[torch.Tensor],
    dst_ptrs: List[torch.Tensor],
    src_indices: torch.Tensor,
    dst_indices: torch.Tensor,
    layer_id: int,
    page_size: int,
):
    torch.ops.sgl_kernel.transfer_kv_per_layer_direct_pf_lf(
        src_ptrs, dst_ptrs, src_indices, dst_indices, layer_id, page_size
    )


def transfer_kv_all_layer_direct_lf_pf(
    src_ptrs: List[torch.Tensor],
    dst_ptrs: List[torch.Tensor],
    src_indices: torch.Tensor,
    dst_indices: torch.Tensor,
    page_size: int,
):
    torch.ops.sgl_kernel.transfer_kv_all_layer_direct_lf_pf(
        src_ptrs, dst_ptrs, src_indices, dst_indices, page_size
172
173
174
    )


175
176
177
178
179
180
181
def transfer_kv_per_layer_mla(
    src: torch.Tensor,
    dst: torch.Tensor,
    src_indices: torch.Tensor,
    dst_indices: torch.Tensor,
    item_size: int,
    block_quota: int = 2,
182
    num_warps_per_block: int = 16 if _is_hip else 32,
183
):
Zhiqiang Xie's avatar
Zhiqiang Xie committed
184
185
186
187
188
189
190
191
192
    torch.ops.sgl_kernel.transfer_kv_per_layer_mla(
        src,
        dst,
        src_indices,
        dst_indices,
        item_size,
        block_quota,
        num_warps_per_block,
    )
193
194


195
def transfer_kv_per_layer_mla_pf_lf(
196
197
198
199
    src: torch.Tensor,
    dst: torch.Tensor,
    src_indices: torch.Tensor,
    dst_indices: torch.Tensor,
200
    layer_id: int,
201
202
203
    item_size: int,
    src_layout_dim: int,
    block_quota: int = 2,
204
    num_warps_per_block: int = 16 if _is_hip else 32,
205
206
207
208
209
210
):
    torch.ops.sgl_kernel.transfer_kv_per_layer_mla_pf_lf(
        src,
        dst,
        src_indices,
        dst_indices,
211
        layer_id,
212
213
214
215
216
217
218
219
220
221
222
223
        item_size,
        src_layout_dim,
        block_quota,
        num_warps_per_block,
    )


def transfer_kv_all_layer_mla(
    src_layers: torch.Tensor,
    dst_layers: torch.Tensor,
    src_indices: torch.Tensor,
    dst_indices: torch.Tensor,
224
225
226
    item_size: int,
    num_layers: int,
    block_quota: int = 2,
227
    num_warps_per_block: int = 16 if _is_hip else 32,
228
):
Zhiqiang Xie's avatar
Zhiqiang Xie committed
229
230
231
232
233
234
235
236
237
238
    torch.ops.sgl_kernel.transfer_kv_all_layer_mla(
        src_layers,
        dst_layers,
        src_indices,
        dst_indices,
        item_size,
        num_layers,
        block_quota,
        num_warps_per_block,
    )
239
240
241
242
243
244
245
246
247
248
249


def transfer_kv_all_layer_mla_lf_pf(
    src_layers: torch.Tensor,
    dst: torch.Tensor,
    src_indices: torch.Tensor,
    dst_indices: torch.Tensor,
    item_size: int,
    dst_layout_dim: int,
    num_layers: int,
    block_quota: int = 2,
250
    num_warps_per_block: int = 16 if _is_hip else 32,
251
252
253
254
255
256
257
258
259
260
261
262
):
    torch.ops.sgl_kernel.transfer_kv_all_layer_mla_lf_pf(
        src_layers,
        dst,
        src_indices,
        dst_indices,
        item_size,
        dst_layout_dim,
        num_layers,
        block_quota,
        num_warps_per_block,
    )