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 3, 09:15 UTC
[Refactor] Remove dead cuda kernels (#42767)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
May 18, 2026 at 14:14 UTC
00e20e76f775b88f47469ae9fcb0f1ecd7580bb9
6 files changed
+1
-605
CMakeLists.txt
-1
@@ -311,7 +311,6 @@ set(VLLM_EXT_SRC
311
"csrc/attention/paged_attention_v1.cu"
312
"csrc/attention/paged_attention_v2.cu"
313
"csrc/attention/merge_attn_states.cu"
314
- "csrc/attention/vertical_slash_index.cu"
314
"csrc/pos_encoding_kernels.cu"
315
"csrc/activation_kernels.cu"
316
"csrc/layernorm_kernels.cu"
csrc/attention/vertical_slash_index.cu
deleted
-401
@@ -1,401 +0,0 @@
1
-// Copyright (c) Microsoft Corporation.
2
-// Licensed under the MIT license.
3
-
4
-#include <assert.h>
5
-
6
-#include <cuda.h>
7
-
8
-#include <torch/all.h>
9
-
10
-__device__ int64_t save_blocks(int* block_offset, int64_t range_start,
11
- int64_t range_end, int64_t block_size,
12
- int64_t input_block_count, int64_t kv_seqlen) {
13
- if (range_start >= kv_seqlen) {
14
- return input_block_count;
15
- }
16
- if (range_end > kv_seqlen) {
17
- range_end = kv_seqlen;
18
- }
19
- int64_t current_block_count = input_block_count;
20
- for (int idx = range_start; idx < range_end; idx += block_size) {
21
- block_offset[current_block_count++] = idx;
22
- }
23
- return current_block_count;
24
-}
25
-
26
-__global__ void convert_vertical_slash_indexes_kernel(
27
- const int* q_seqlens, // [BATCH, ]
28
- const int* kv_seqlens, // [BATCH, ]
29
- const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
30
- const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S]
31
- int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
32
- int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
33
- int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
34
- int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
35
- int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
36
- int64_t NNZ_V, int64_t NNZ_S,
37
- bool causal // True for intra, False for succ
38
-) {
39
- const int batch_idx = blockIdx.y;
40
- const int head_idx = blockIdx.x;
41
- const int group_idx = blockIdx.z;
42
-
43
- int64_t q_seqlen = q_seqlens[batch_idx];
44
- int64_t kv_seqlen = kv_seqlens[batch_idx];
45
- int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x;
46
- int64_t start_m = block_idx_m * BLOCK_SIZE_M;
47
- if (start_m >= q_seqlen) {
48
- return;
49
- }
50
- int64_t end_m = start_m + BLOCK_SIZE_M;
51
- vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V;
52
- slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S;
53
- int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m;
54
- block_count += row_offset;
55
- block_offset += row_offset * NNZ_S;
56
- column_count += row_offset;
57
- column_index += row_offset * NNZ_V;
58
-
59
- bool has_slash = true;
60
- int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0;
61
- int64_t s = 0, v = 0;
62
- int64_t v_idx = vertical_indexes[v++];
63
- int64_t s_idx = slash_indexes[s++];
64
- if (causal) {
65
- while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) {
66
- s_idx = slash_indexes[s++];
67
- }
68
- if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false;
69
- s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M);
70
- } else {
71
- while (s_idx >= end_m + kv_seqlen && s < NNZ_S) {
72
- s_idx = slash_indexes[s++];
73
- }
74
- if (s_idx > end_m + kv_seqlen) has_slash = false;
75
- s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M);
76
- }
77
-
78
- int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx;
79
- if (!has_slash) {
80
- if (causal) {
81
- range_start = (kv_seqlen - q_seqlen) + end_m;
82
- range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
83
- } else {
84
- range_start = kv_seqlen;
85
- range_end = kv_seqlen + BLOCK_SIZE_N;
86
- }
87
- }
88
-
89
- bool slash_finished = false;
90
- while (1) {
91
- if (v_idx < range_end) {
92
- if (v_idx < range_start) {
93
- column_index[tmp_col_cnt++] = v_idx;
94
- }
95
- if (v < NNZ_V) {
96
- v_idx = vertical_indexes[v++];
97
- } else {
98
- if (causal)
99
- v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen);
100
- else
101
- v_idx = end_m + BLOCK_SIZE_N + kv_seqlen;
102
- }
103
- } else {
104
- if ((s < NNZ_S && causal) ||
105
- (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) {
106
- if (causal)
107
- s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++],
108
- BLOCK_SIZE_M);
109
- else
110
- s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M);
111
- } else {
112
- if (v == NNZ_V || (v_idx > range_start && causal)) {
113
- // add the last vertical if no more slash
114
- if (v == NNZ_V && !causal && v_idx < kv_seqlen) {
115
- column_index[tmp_col_cnt++] = v_idx;
116
- }
117
- tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
118
- BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
119
- break;
120
- } else {
121
- if (causal) {
122
- range_start = (kv_seqlen - q_seqlen) + end_m;
123
- range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
124
- } else {
125
- // if slash_finished but there are vertical left, save current
126
- // blocks
127
- tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
128
- BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
129
- range_start = kv_seqlen;
130
- range_end = kv_seqlen + BLOCK_SIZE_N;
131
- }
132
- slash_finished = true;
133
- }
134
- }
135
- if (!slash_finished) {
136
- if (s_idx > range_end + BLOCK_SIZE_M) {
137
- tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
138
- BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
139
- range_start = s_idx - BLOCK_SIZE_M;
140
- range_end = s_idx;
141
- } else if (s_idx > range_end) {
142
- range_end += BLOCK_SIZE_M;
143
- }
144
- }
145
- }
146
- }
147
-
148
- block_count[0] = tmp_blk_cnt;
149
- column_count[0] = tmp_col_cnt;
150
-}
151
-
152
-void convert_vertical_slash_indexes_64x64(
153
- const int* q_seqlens, // [BATCH, ]
154
- const int* kv_seqlens, // [BATCH, ]
155
- const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
156
- const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S]
157
- int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
158
- int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
159
- int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
160
- int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
161
- int64_t BATCH_SIZE, int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M,
162
- int64_t BLOCK_SIZE_N, int64_t NNZ_V, int64_t NNZ_S, bool causal) {
163
- const int N_THREADS = 64;
164
- const dim3 dimBlock(N_THREADS);
165
- const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS);
166
- convert_vertical_slash_indexes_kernel<<<dimGrid, dimBlock>>>(
167
- q_seqlens, kv_seqlens, vertical_indexes, slash_indexes, block_count,
168
- block_offset, column_count, column_index, N_HEADS, N_ROWS, BLOCK_SIZE_M,
169
- BLOCK_SIZE_N, NNZ_V, NNZ_S, causal);
170
-}
171
-
172
-/**
173
- * Implements the Algorithm 4 in paper https://arxiv.org/abs/2407.02490.
174
- *
175
- * This function builds the index of each row of blocks from vertical indices
176
- * and slash indices. The vertical indices are treated as points, while the
177
- * slash indices are converted as ranges. The output consists of the merged
178
- * ranges and separate column indices, where the ranges are represented by
179
- * block indices.
180
- *
181
- * The implementation is referenced from the original MInference repo:
182
- * https://github.com/microsoft/MInference/blob/main/csrc/vertical_slash_index.cu.
183
- */
184
-void convert_vertical_slash_indexes(
185
- torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]
186
- torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
187
- torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS]
188
- torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
189
- torch::Tensor q_seqlens, // [BATCH, ]
190
- torch::Tensor kv_seqlens, // [BATCH, ]
191
- torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
192
- torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S]
193
- int64_t context_size, int64_t block_size_M, int64_t block_size_N,
194
- bool causal) {
195
- cudaSetDevice(q_seqlens.get_device());
196
-
197
- int batch_size = slash_indexes.size(0);
198
- int num_heads = slash_indexes.size(1);
199
- int nnz_slash = slash_indexes.size(2);
200
- int nnz_vertical = vertical_indexes.size(2);
201
- int num_rows = (context_size + block_size_M - 1) / block_size_M;
202
-
203
- convert_vertical_slash_indexes_64x64(
204
- q_seqlens.data_ptr<int>(), kv_seqlens.data_ptr<int>(),
205
- vertical_indexes.data_ptr<int>(), slash_indexes.data_ptr<int>(),
206
- block_count.data_ptr<int>(), block_offset.data_ptr<int>(),
207
- column_count.data_ptr<int>(), column_index.data_ptr<int>(), batch_size,
208
- num_heads, num_rows, block_size_M, block_size_N, nnz_vertical, nnz_slash,
209
- causal);
210
-}
211
-
212
-__global__ void convert_vertical_slash_indexes_kernel_mergehead(
213
- const int* q_seqlens, // [BATCH, ]
214
- const int* kv_seqlens, // [BATCH, ]
215
- const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
216
- const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S]
217
- const int* per_head_vertical_topkv, const int* per_head_slash_topkv,
218
- int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
219
- int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
220
- int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
221
- int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
222
- int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
223
- int64_t NNZ_V, int64_t NNZ_S,
224
- bool causal // True for intra, False for succ
225
-) {
226
- const int batch_idx = blockIdx.y;
227
- const int head_idx = blockIdx.x;
228
- const int group_idx = blockIdx.z;
229
-
230
- int64_t q_seqlen = q_seqlens[batch_idx];
231
- int64_t kv_seqlen = kv_seqlens[batch_idx];
232
- int64_t block_idx_m = group_idx * blockDim.x + threadIdx.x;
233
- int64_t start_m = block_idx_m * BLOCK_SIZE_M;
234
- if (start_m >= q_seqlen) {
235
- return;
236
- }
237
- int64_t end_m = start_m + BLOCK_SIZE_M;
238
- vertical_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_V;
239
- slash_indexes += (batch_idx * N_HEADS + head_idx) * NNZ_S;
240
- int64_t row_offset = (batch_idx * N_HEADS + head_idx) * N_ROWS + block_idx_m;
241
- block_count += row_offset;
242
- block_offset += row_offset * NNZ_S;
243
- column_count += row_offset;
244
- column_index += row_offset * NNZ_V;
245
-
246
- // MergeHead: each head has it's unique max topk NNZ_V，NNZ_S. (NNZ_V，NNZ_S
247
- // above is buffer size, use to compute offset)
248
- NNZ_S = per_head_slash_topkv[head_idx];
249
- NNZ_V = per_head_vertical_topkv[head_idx];
250
-
251
- bool has_slash = true;
252
- int64_t tmp_col_cnt = 0, tmp_blk_cnt = 0;
253
- int64_t s = 0, v = 0;
254
- int64_t v_idx = vertical_indexes[v++];
255
- int64_t s_idx = slash_indexes[s++];
256
- if (causal) {
257
- while (s_idx >= end_m + (kv_seqlen - q_seqlen) && s < NNZ_S) {
258
- s_idx = slash_indexes[s++];
259
- }
260
- if (s_idx > end_m + (kv_seqlen - q_seqlen)) has_slash = false;
261
- s_idx = max((kv_seqlen - q_seqlen) + end_m - s_idx, BLOCK_SIZE_M);
262
- } else {
263
- while (s_idx >= end_m + kv_seqlen && s < NNZ_S) {
264
- s_idx = slash_indexes[s++];
265
- }
266
- if (s_idx > end_m + kv_seqlen) has_slash = false;
267
- s_idx = max(kv_seqlen + end_m - s_idx, BLOCK_SIZE_M);
268
- }
269
-
270
- int64_t range_start = s_idx - BLOCK_SIZE_M, range_end = s_idx;
271
- if (!has_slash) {
272
- if (causal) {
273
- range_start = (kv_seqlen - q_seqlen) + end_m;
274
- range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
275
- } else {
276
- range_start = kv_seqlen;
277
- range_end = kv_seqlen + BLOCK_SIZE_N;
278
- }
279
- }
280
-
281
- bool slash_finished = false;
282
- while (1) {
283
- if (v_idx < range_end) {
284
- if (v_idx < range_start) {
285
- column_index[tmp_col_cnt++] = v_idx;
286
- }
287
- if (v < NNZ_V) {
288
- v_idx = vertical_indexes[v++];
289
- } else {
290
- if (causal)
291
- v_idx = end_m + BLOCK_SIZE_N + (kv_seqlen - q_seqlen);
292
- else
293
- v_idx = end_m + BLOCK_SIZE_N + kv_seqlen;
294
- }
295
- } else {
296
- if ((s < NNZ_S && causal) ||
297
- (s < NNZ_S && !causal && slash_indexes[s] >= start_m)) {
298
- if (causal)
299
- s_idx = max((kv_seqlen - q_seqlen) + end_m - slash_indexes[s++],
300
- BLOCK_SIZE_M);
301
- else
302
- s_idx = max(kv_seqlen + end_m - slash_indexes[s++], BLOCK_SIZE_M);
303
- } else {
304
- if (v == NNZ_V || (v_idx > range_start && causal)) {
305
- // add the last vertical if no more slash
306
- if (v == NNZ_V && !causal && v_idx < kv_seqlen) {
307
- column_index[tmp_col_cnt++] = v_idx;
308
- }
309
- tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
310
- BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
311
- break;
312
- } else {
313
- if (causal) {
314
- range_start = (kv_seqlen - q_seqlen) + end_m;
315
- range_end = (kv_seqlen - q_seqlen) + end_m + BLOCK_SIZE_N;
316
- } else {
317
- // if slash_finished but there are vertical left, save current
318
- // blocks
319
- tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
320
- BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
321
- range_start = kv_seqlen;
322
- range_end = kv_seqlen + BLOCK_SIZE_N;
323
- }
324
- slash_finished = true;
325
- }
326
- }
327
- if (!slash_finished) {
328
- if (s_idx > range_end + BLOCK_SIZE_M) {
329
- tmp_blk_cnt = save_blocks(block_offset, range_start, range_end,
330
- BLOCK_SIZE_N, tmp_blk_cnt, kv_seqlen);
331
- range_start = s_idx - BLOCK_SIZE_M;
332
- range_end = s_idx;
333
- } else if (s_idx > range_end) {
334
- range_end += BLOCK_SIZE_M;
335
- }
336
- }
337
- }
338
- }
339
-
340
- block_count[0] = tmp_blk_cnt;
341
- column_count[0] = tmp_col_cnt;
342
-}
343
-
344
-void convert_vertical_slash_indexes_64x64_mergehead(
345
- const int* q_seqlens, // [BATCH, ]
346
- const int* kv_seqlens, // [BATCH, ]
347
- const int* vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
348
- const int* slash_indexes, // [BATCH, N_HEADS, NNZ_S]
349
- int* per_head_vertical_topkv, int* per_head_slash_topkv,
350
- int* block_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
351
- int* block_offset, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_S]
352
- int* column_count, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M)]
353
- int* column_index, // [BATCH, N_HEADS, cdiv(N_CTX, BLOCK_SIZE_M), NNZ_V]
354
- int64_t BATCH_SIZE, int64_t N_HEADS, int64_t N_ROWS, int64_t BLOCK_SIZE_M,
355
- int64_t BLOCK_SIZE_N, int64_t NNZ_V, int64_t NNZ_S, bool causal) {
356
- const int N_THREADS = 64;
357
- const dim3 dimBlock(N_THREADS);
358
- const dim3 dimGrid(N_HEADS, BATCH_SIZE, (N_ROWS + N_THREADS - 1) / N_THREADS);
359
- convert_vertical_slash_indexes_kernel_mergehead<<<dimGrid, dimBlock>>>(
360
- q_seqlens, kv_seqlens, vertical_indexes, slash_indexes,
361
- per_head_vertical_topkv, per_head_slash_topkv, block_count, block_offset,
362
- column_count, column_index, N_HEADS, N_ROWS, BLOCK_SIZE_M, BLOCK_SIZE_N,
363
- NNZ_V, NNZ_S, causal);
364
-}
365
-
366
-/**
367
- * Implements the Algorithm 4 in paper https://arxiv.org/abs/2407.02490.
368
- *
369
- * Like the above convert_vertical_slash_indexes, but with
370
- * pre-computed vertical and slash counts.
371
- */
372
-void convert_vertical_slash_indexes_mergehead(
373
- torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]
374
- torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
375
- torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS]
376
- torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
377
- torch::Tensor q_seqlens, // [BATCH, ]
378
- torch::Tensor kv_seqlens, // [BATCH, ]
379
- torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
380
- torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S]
381
- torch::Tensor vertical_indices_count, // [N_HEADS, ]
382
- torch::Tensor slash_indices_count, // [N_HEADS, ]
383
- int64_t context_size, int64_t block_size_M, int64_t block_size_N,
384
- bool causal) {
385
- cudaSetDevice(q_seqlens.get_device());
386
-
387
- int batch_size = slash_indexes.size(0);
388
- int num_heads = slash_indexes.size(1);
389
- int nnz_slash = slash_indexes.size(2);
390
- int nnz_vertical = vertical_indexes.size(2);
391
- int num_rows = (context_size + block_size_M - 1) / block_size_M;
392
-
393
- convert_vertical_slash_indexes_64x64_mergehead(
394
- q_seqlens.data_ptr<int>(), kv_seqlens.data_ptr<int>(),
395
- vertical_indexes.data_ptr<int>(), slash_indexes.data_ptr<int>(),
396
- vertical_indices_count.data_ptr<int>(),
397
- slash_indices_count.data_ptr<int>(), block_count.data_ptr<int>(),
398
- block_offset.data_ptr<int>(), column_count.data_ptr<int>(),
399
- column_index.data_ptr<int>(), batch_size, num_heads, num_rows,
400
- block_size_M, block_size_N, nnz_vertical, nnz_slash, causal);
401
-}
csrc/moe/torch_bindings.cpp
-10
@@ -92,16 +92,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
92
"bool use_fp32_reduce, bool is_zp_float,"
93
"int thread_k, int thread_n, int blocks_per_sm) -> Tensor");
94
95
- m.def(
96
- "marlin_gemm_moe(Tensor! a, Tensor! b_q_weights, Tensor! sorted_ids, "
97
- "Tensor! topk_weights, Tensor! topk_ids, Tensor! b_scales, Tensor! "
98
- "b_zeros, Tensor! g_idx, Tensor! perm, Tensor! workspace, "
99
- "int b_q_type, SymInt size_m, "
100
- "SymInt size_n, SymInt size_k, bool is_k_full, int num_experts, int "
101
- "topk, "
102
- "int moe_block_size, bool replicate_input, bool apply_weights)"
103
- " -> Tensor");
104
-
95
m.def(
96
"moe_permute(Tensor input, Tensor topk_ids,"
97
"Tensor token_expert_indices, Tensor? expert_map, int n_expert,"
csrc/ops.h
-26
@@ -60,32 +60,6 @@ void merge_attn_states(
60
const torch::Tensor& suffix_output, const torch::Tensor& suffix_lse,
61
const std::optional<int64_t> prefill_tokens_with_context,
62
const std::optional<torch::Tensor>& output_scale = std::nullopt);
63
-#ifndef USE_ROCM
64
-void convert_vertical_slash_indexes(
65
- torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]
66
- torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
67
- torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS]
68
- torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
69
- torch::Tensor q_seqlens, // [BATCH, ]
70
- torch::Tensor kv_seqlens, // [BATCH, ]
71
- torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
72
- torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S]
73
- int64_t context_size, int64_t block_size_M, int64_t block_size_N,
74
- bool causal);
75
-
76
-void convert_vertical_slash_indexes_mergehead(
77
- torch::Tensor& block_count, // [BATCH, N_HEADS, NUM_ROWS]
78
- torch::Tensor& block_offset, // [BATCH, N_HEADS, NUM_ROWS, NNZ_S]
79
- torch::Tensor& column_count, // [BATCH, N_HEADS, NUM_ROWS]
80
- torch::Tensor& column_index, // [BATCH, N_HEADS, NUM_ROWS, NNZ_V]
81
- torch::Tensor q_seqlens, // [BATCH, ]
82
- torch::Tensor kv_seqlens, // [BATCH, ]
83
- torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
84
- torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S]
85
- torch::Tensor vertical_indices_count, // [N_HEADS, ]
86
- torch::Tensor slash_indices_count, int64_t context_size,
87
- int64_t block_size_M, int64_t block_size_N, bool causal);
88
-#endif
63
64
void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
65
double epsilon);
csrc/torch_bindings.cpp
-24
@@ -76,30 +76,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
76
" int!? prefill_tokens_with_context,"
77
" Tensor? output_scale=None) -> ()");
78
ops.impl("merge_attn_states", torch::kCUDA, &merge_attn_states);
79
-#ifndef USE_ROCM
80
- ops.def(
81
- "convert_vertical_slash_indexes("
82
- " Tensor! block_count, Tensor! block_offset, "
83
- " Tensor! column_count, Tensor! column_index, "
84
- " Tensor q_seqlens, Tensor q_seqlens, "
85
- " Tensor vertical_indexes, Tensor slash_indexes, "
86
- " int context_size, int block_size_M, int block_size_N, "
87
- " bool causal) -> ()");
88
- ops.impl("convert_vertical_slash_indexes", torch::kCUDA,
89
- &convert_vertical_slash_indexes);
90
-
91
- ops.def(
92
- "convert_vertical_slash_indexes_mergehead("
93
- " Tensor! block_count, Tensor! block_offset, "
94
- " Tensor! column_count, Tensor! column_index, "
95
- " Tensor q_seqlens, Tensor q_seqlens, "
96
- " Tensor vertical_indexes, Tensor slash_indexes, "
97
- " Tensor vertical_indices_count, Tensor slash_indices_count, "
98
- " int context_size, int block_size_M, int block_size_N, "
99
- " bool causal) -> ()");
100
- ops.impl("convert_vertical_slash_indexes_mergehead", torch::kCUDA,
101
- &convert_vertical_slash_indexes_mergehead);
102
-#endif
79
80
// Activation ops
81
// Activation function used in SwiGLU.
vllm/_custom_ops.py
+1
-143
@@ -280,123 +280,6 @@ def merge_attn_states(
280
)
281
282
283
-def convert_vertical_slash_indexes(
284
- q_seqlens: torch.Tensor, # [BATCH, ]
285
- kv_seqlens: torch.Tensor, # [BATCH, ]
286
- vertical_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_V]
287
- slash_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_S]
288
- context_size: int,
289
- block_size_M: int,
290
- block_size_N: int,
291
- causal: bool = True,
292
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
293
- batch_size = slash_indexes.size(0)
294
- num_heads = slash_indexes.size(1)
295
- nnz_slash = slash_indexes.size(2)
296
- nnz_vertical = vertical_indexes.size(2)
297
- num_rows = (context_size + block_size_M - 1) // block_size_M
298
-
299
- block_count = torch.zeros(
300
- batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device
301
- )
302
- block_offset = torch.zeros(
303
- batch_size,
304
- num_heads,
305
- num_rows,
306
- nnz_slash,
307
- dtype=q_seqlens.dtype,
308
- device=q_seqlens.device,
309
- )
310
- column_count = torch.zeros(
311
- batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device
312
- )
313
- column_index = torch.zeros(
314
- batch_size,
315
- num_heads,
316
- num_rows,
317
- nnz_vertical,
318
- dtype=q_seqlens.dtype,
319
- device=q_seqlens.device,
320
- )
321
-
322
- torch.ops._C.convert_vertical_slash_indexes(
323
- block_count,
324
- block_offset,
325
- column_count,
326
- column_index,
327
- q_seqlens,
328
- kv_seqlens,
329
- vertical_indexes,
330
- slash_indexes,
331
- context_size,
332
- block_size_M,
333
- block_size_N,
334
- causal,
335
- )
336
- return block_count, block_offset, column_count, column_index
337
-
338
-
339
-def convert_vertical_slash_indexes_mergehead(
340
- q_seqlens: torch.Tensor, # [BATCH, ]
341
- kv_seqlens: torch.Tensor, # [BATCH, ]
342
- vertical_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_V]
343
- slash_indexes: torch.Tensor, # [BATCH, N_HEADS, NNZ_S]
344
- # [N_HEADS] : different head use different number of indices
345
- vertical_indices_count: torch.Tensor,
346
- slash_indices_count: torch.Tensor,
347
- context_size: int,
348
- block_size_M: int,
349
- block_size_N: int,
350
- causal: bool = True,
351
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
352
- batch_size = slash_indexes.size(0)
353
- num_heads = slash_indexes.size(1)
354
- nnz_slash = slash_indexes.size(2)
355
- nnz_vertical = vertical_indexes.size(2)
356
- num_rows = (context_size + block_size_M - 1) // block_size_M
357
-
358
- block_count = torch.empty(
359
- batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device
360
- )
361
- block_offset = torch.empty(
362
- batch_size,
363
- num_heads,
364
- num_rows,
365
- nnz_slash,
366
- dtype=q_seqlens.dtype,
367
- device=q_seqlens.device,
368
- )
369
- column_count = torch.empty(
370
- batch_size, num_heads, num_rows, dtype=q_seqlens.dtype, device=q_seqlens.device
371
- )
372
- column_index = torch.empty(
373
- batch_size,
374
- num_heads,
375
- num_rows,
376
- nnz_vertical,
377
- dtype=q_seqlens.dtype,
378
- device=q_seqlens.device,
379
- )
380
-
381
- torch.ops._C.convert_vertical_slash_indexes_mergehead(
382
- block_count,
383
- block_offset,
384
- column_count,
385
- column_index,
386
- q_seqlens,
387
- kv_seqlens,
388
- vertical_indexes,
389
- slash_indexes,
390
- vertical_indices_count,
391
- slash_indices_count,
392
- context_size,
393
- block_size_M,
394
- block_size_N,
395
- causal,
396
- )
397
- return block_count, block_offset, column_count, column_index
398
-
399
-
283
# pos encoding ops
284
def rotary_embedding(
285
positions: torch.Tensor,
@@ -2658,32 +2541,7 @@ def moe_wna16_marlin_gemm(
2541
)
2542
2543
2661
-if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "marlin_gemm_moe"):
2662
-
2663
- @register_fake("_moe_C::marlin_gemm_moe")
2664
- def marlin_gemm_moe_fake(
2665
- a: torch.Tensor,
2666
- b_q_weights: torch.Tensor,
2667
- sorted_ids: torch.Tensor,
2668
- topk_weights: torch.Tensor,
2669
- topk_ids: torch.Tensor,
2670
- b_scales: torch.Tensor,
2671
- b_zero_points: torch.Tensor,
2672
- g_idx: torch.Tensor,
2673
- perm: torch.Tensor,
2674
- workspace: torch.Tensor,
2675
- b_q_type: ScalarType,
2676
- size_m: torch.SymInt,
2677
- size_n: torch.SymInt,
2678
- size_k: torch.SymInt,
2679
- is_k_full: bool,
2680
- num_experts: int,
2681
- topk: int,
2682
- moe_block_size: int,
2683
- replicate_input: bool,
2684
- apply_weights: bool,
2685
- ) -> torch.Tensor:
2686
- return torch.empty((size_m, topk, size_n), dtype=a.dtype, device=a.device)
2544
+if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "moe_wna16_marlin_gemm"):
2545
2546
@register_fake("_moe_C::moe_wna16_marlin_gemm")
2547
def moe_wna16_marlin_gemm_fake(