_custom_ops.py 15.8 KB
Newer Older
1
2
import contextlib
from typing import List, Optional, Tuple, Type
3
4
5
6

import torch

try:
7
    import vllm._C
8
9
10
11
except ImportError as e:
    from vllm.logger import init_logger
    logger = init_logger(__name__)
    logger.warning("Failed to import from vllm._C with %r", e)
12

13
14
15
16
17
18
19
20
21
22
23
with contextlib.suppress(ImportError):
    import vllm._moe_C

with contextlib.suppress(ImportError):
    # ruff: noqa: F401
    import vllm._punica_C


def is_custom_op_supported(op_name: str) -> bool:
    op, overloads = torch._C._jit_get_operation(op_name)
    return op is not None
24
25
26
27


# activation ops
def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
28
    torch.ops._C.silu_and_mul(out, x)
29
30
31


def gelu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
32
    torch.ops._C.gelu_and_mul(out, x)
33
34
35


def gelu_tanh_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
36
    torch.ops._C.gelu_tanh_and_mul(out, x)
37
38
39


def gelu_fast(out: torch.Tensor, x: torch.Tensor) -> None:
40
    torch.ops._C.gelu_fast(out, x)
41
42
43


def gelu_new(out: torch.Tensor, x: torch.Tensor) -> None:
44
    torch.ops._C.gelu_new(out, x)
45
46
47
48
49
50
51
52
53
54
55


# page attention ops
def paged_attention_v1(
    out: torch.Tensor,
    query: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    num_kv_heads: int,
    scale: float,
    block_tables: torch.Tensor,
56
    seq_lens: torch.Tensor,
57
    block_size: int,
58
    max_seq_len: int,
59
60
61
    alibi_slopes: Optional[torch.Tensor],
    kv_cache_dtype: str,
    kv_scale: float,
62
63
64
65
66
    tp_rank: int = 0,
    blocksparse_local_blocks: int = 0,
    blocksparse_vert_stride: int = 0,
    blocksparse_block_size: int = 64,
    blocksparse_head_sliding_step: int = 0,
67
) -> None:
68
    torch.ops._C.paged_attention_v1(
69
70
71
72
        out, query, key_cache, value_cache, num_kv_heads, scale, block_tables,
        seq_lens, block_size, max_seq_len, alibi_slopes, kv_cache_dtype,
        kv_scale, tp_rank, blocksparse_local_blocks, blocksparse_vert_stride,
        blocksparse_block_size, blocksparse_head_sliding_step)
73
74
75
76
77
78
79
80
81
82
83
84
85


def paged_attention_v2(
    out: torch.Tensor,
    exp_sum: torch.Tensor,
    max_logits: torch.Tensor,
    tmp_out: torch.Tensor,
    query: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    num_kv_heads: int,
    scale: float,
    block_tables: torch.Tensor,
86
    seq_lens: torch.Tensor,
87
    block_size: int,
88
    max_seq_len: int,
89
90
91
    alibi_slopes: Optional[torch.Tensor],
    kv_cache_dtype: str,
    kv_scale: float,
92
93
94
95
96
    tp_rank: int = 0,
    blocksparse_local_blocks: int = 0,
    blocksparse_vert_stride: int = 0,
    blocksparse_block_size: int = 64,
    blocksparse_head_sliding_step: int = 0,
97
) -> None:
98
    torch.ops._C.paged_attention_v2(
99
100
101
102
103
        out, exp_sum, max_logits, tmp_out, query, key_cache, value_cache,
        num_kv_heads, scale, block_tables, seq_lens, block_size, max_seq_len,
        alibi_slopes, kv_cache_dtype, kv_scale, tp_rank,
        blocksparse_local_blocks, blocksparse_vert_stride,
        blocksparse_block_size, blocksparse_head_sliding_step)
104
105
106
107
108
109
110
111
112
113
114


# pos encoding ops
def rotary_embedding(
    positions: torch.Tensor,
    query: torch.Tensor,
    key: torch.Tensor,
    head_size: int,
    cos_sin_cache: torch.Tensor,
    is_neox: bool,
) -> None:
115
116
    torch.ops._C.rotary_embedding(positions, query, key, head_size,
                                  cos_sin_cache, is_neox)
117
118
119
120
121
122
123


def batched_rotary_embedding(positions: torch.Tensor, query: torch.Tensor,
                             key: torch.Tensor, head_size: int,
                             cos_sin_cache: torch.Tensor, is_neox: bool,
                             rot_dim: int,
                             cos_sin_cache_offsets: torch.Tensor) -> None:
124
125
126
    torch.ops._C.batched_rotary_embedding(positions, query, key, head_size,
                                          cos_sin_cache, is_neox, rot_dim,
                                          cos_sin_cache_offsets)
127
128
129
130
131


# layer norm ops
def rms_norm(out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor,
             epsilon: float) -> None:
132
    torch.ops._C.rms_norm(out, input, weight, epsilon)
133
134
135
136


def fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor,
                       weight: torch.Tensor, epsilon: float) -> None:
137
    torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon)
138
139
140
141
142
143
144


# quantization ops
# awq
def awq_dequantize(qweight: torch.Tensor, scales: torch.Tensor,
                   zeros: torch.Tensor, split_k_iters: int, thx: int,
                   thy: int) -> torch.Tensor:
145
146
    return torch.ops._C.awq_dequantize(qweight, scales, zeros, split_k_iters,
                                       thx, thy)
147
148
149
150


def awq_gemm(input: torch.Tensor, qweight: torch.Tensor, qzeros: torch.Tensor,
             scales: torch.Tensor, split_k_iters: int) -> torch.Tensor:
151
    return torch.ops._C.awq_gemm(input, qweight, qzeros, scales, split_k_iters)
152
153
154
155
156
157
158


# gptq
def gptq_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
              b_gptq_qzeros: torch.Tensor, b_gptq_scales: torch.Tensor,
              b_g_idx: torch.Tensor, use_exllama: bool,
              bit: int) -> torch.Tensor:
159
160
    return torch.ops._C.gptq_gemm(a, b_q_weight, b_gptq_qzeros, b_gptq_scales,
                                  b_g_idx, use_exllama, bit)
161
162
163
164


def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor,
                 bit: int) -> None:
165
    torch.ops._C.gptq_shuffle(q_weight, q_perm, bit)
166

gaoqiong's avatar
gaoqiong committed
167
168
169
170
# trans_w16
def trans_w16_gemm(dst: torch.Tensor, src: torch.Tensor,
                row:int, col:int) -> None :
    torch.ops._C.trans_w16_gemm(dst,src,row,col)
171
172
173
174

# squeezellm
def squeezellm_gemm(vec: torch.Tensor, mat: torch.Tensor, mul: torch.Tensor,
                    lookup_table: torch.Tensor) -> None:
175
    torch.ops._C.squeezellm_gemm(vec, mat, mul, lookup_table)
176
177
178
179
180
181


# marlin
def marlin_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
                b_scales: torch.Tensor, workspace: torch.Tensor, size_m: int,
                size_n: int, size_k: int) -> torch.Tensor:
182
183
    return torch.ops._C.marlin_gemm(a, b_q_weight, b_scales, workspace, size_m,
                                    size_n, size_k)
184
185


186
187
188
189
190
# marlin_24
def gptq_marlin_24_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
                        b_meta: torch.Tensor, b_scales: torch.Tensor,
                        workspace: torch.Tensor, num_bits: int, size_m: int,
                        size_n: int, size_k: int) -> torch.Tensor:
191
192
193
    return torch.ops._C.gptq_marlin_24_gemm(a, b_q_weight, b_meta, b_scales,
                                            workspace, num_bits, size_m,
                                            size_n, size_k)
194
195


196
197
# cutlass
def cutlass_scaled_mm_dq(a: torch.Tensor, b: torch.Tensor,
198
                         scale_a: torch.Tensor, scale_b: torch.Tensor,
199
200
201
202
203
204
205
206
                         out_dtype: Type[torch.dtype]) -> torch.Tensor:
    assert (b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0)
    assert (out_dtype is torch.bfloat16 or out_dtype is torch.float16)

    m = a.shape[0]
    n = b.shape[1]
    out = torch.empty((m, n), dtype=out_dtype, device=a.device)

207
    torch.ops._C.cutlass_scaled_mm_dq(out, a, b, scale_a, scale_b)
208
209
210
211

    return out


212
213
214
215
216
# aqlm
def aqlm_gemm(input: torch.Tensor, codes: torch.Tensor,
              codebooks: torch.Tensor, scales: torch.Tensor,
              codebook_partition_sizes: torch.Tensor,
              bias: Optional[torch.Tensor]) -> torch.Tensor:
217
218
    return torch.ops._C.aqlm_gemm(input, codes, codebooks, scales,
                                  codebook_partition_sizes, bias)
219
220
221
222


def aqlm_dequant(codes: torch.Tensor, codebooks: torch.Tensor,
                 codebook_partition_sizes: torch.Tensor) -> torch.Tensor:
223
224
    return torch.ops._C.aqlm_dequant(codes, codebooks,
                                     codebook_partition_sizes)
225
226


227
228
# gptq_marlin
def gptq_marlin_repack(b_q_weight: torch.Tensor, perm: torch.Tensor,
229
230
                       size_k: int, size_n: int,
                       num_bits: int) -> torch.Tensor:
231
232
    return torch.ops._C.gptq_marlin_repack(b_q_weight, perm, size_k, size_n,
                                           num_bits)
233
234
235
236


def gptq_marlin_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
                     b_scales: torch.Tensor, g_idx: torch.Tensor,
237
238
                     perm: torch.Tensor, workspace: torch.Tensor,
                     num_bits: int, size_m: int, size_n: int, size_k: int,
239
                     is_k_full: bool) -> torch.Tensor:
240
241
242
    return torch.ops._C.gptq_marlin_gemm(a, b_q_weight, b_scales, g_idx, perm,
                                         workspace, num_bits, size_m, size_n,
                                         size_k, is_k_full)
243
244


245
# fp8
zhuwenwen's avatar
zhuwenwen committed
246
247
248
# def scaled_fp8_quant(
#     input: torch.Tensor,
#     scale: Optional[torch.Tensor] = None,
zhuwenwen's avatar
zhuwenwen committed
249
#     batch_dim_padding: Optional[int] = None,
zhuwenwen's avatar
zhuwenwen committed
250
# ) -> Tuple[torch.Tensor, torch.Tensor]:
zhuwenwen's avatar
zhuwenwen committed
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
#     """
#     Quantize input tensor to FP8 and return quantized tensor and scale.

#     This function supports both static and dynamic quantization: If you
#     provide the scale, it will use static scaling and if you omit it,
#     the scale will be determined dynamically. The function also allows
#     optional padding of the output tensor for downstream kernels that
#     will benefit from padding.

#     Args:
#         input: The input tensor to be quantized to FP8
#         scale: Optional scaling factor for the FP8 quantization
#         batch_dim_padding: If specified, pad the first dimension
#             of the output to at least this value.

#     Returns:
#         Tuple[torch.Tensor, torch.Tensor]: The output tensor in FP8 and
#             scaling factor.
#     """
#     if batch_dim_padding:
#         shape = (max(batch_dim_padding, input.shape[0]), *input.shape[1:])
#         output = torch.empty(shape,
#                              device=input.device,
#                              dtype=torch.float8_e4m3fn)
#     else:
#         output = torch.empty_like(input, dtype=torch.float8_e4m3fn)
zhuwenwen's avatar
zhuwenwen committed
277
278
#     if scale is None:
#         scale = torch.zeros(1, device=input.device, dtype=torch.float32)
zhuwenwen's avatar
zhuwenwen committed
279
#         torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)
zhuwenwen's avatar
zhuwenwen committed
280
#     else:
zhuwenwen's avatar
zhuwenwen committed
281
#         torch.ops._C.static_scaled_fp8_quant(output, input, scale)
zhuwenwen's avatar
zhuwenwen committed
282
#     return output, scale
283
284


285
# int8
286
287
288
289
def scaled_int8_quant(
        input: torch.Tensor,
        scale: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
zhuwenwen's avatar
zhuwenwen committed
290
    """
291
    Quantize the input tensor to int8 and return the quantized tensor and scale.
zhuwenwen's avatar
zhuwenwen committed
292
293
294

    Args:
        input: The input tensor to be quantized to int8.
295
296
        scale: Optional scaling factor for the int8 quantization.
            When not provided, we invoke dynamic-per-token quantization.
zhuwenwen's avatar
zhuwenwen committed
297
298

    Returns:
299
      Tuple[Torch.Tensor, Torch.Tensor] : Output int8 tensor and scales.
zhuwenwen's avatar
zhuwenwen committed
300
    """
301
302
303
    output = torch.empty_like(input, dtype=torch.int8)
    if scale is not None:
        # static-per-tensor quantization.
304
        torch.ops._C.static_scaled_int8_quant(output, input, scale)
305
306
307
308
309
310
        return output, scale

    # dynamic-per-token quantization.
    input_scales = torch.empty((input.numel() // input.shape[-1], 1),
                               device=input.device,
                               dtype=torch.float32)
311
    torch.ops._C.dynamic_scaled_int8_quant(output, input, input_scales)
312
    return output, input_scales
313
314


315
316
317
318
319
# moe
def moe_align_block_size(topk_ids: torch.Tensor, num_experts: int,
                         block_size: int, sorted_token_ids: torch.Tensor,
                         experts_ids: torch.Tensor,
                         num_tokens_post_pad: torch.Tensor) -> None:
320
321
322
323
324
325
326
327
328
329
    torch.ops._C.moe_align_block_size(topk_ids, num_experts, block_size,
                                      sorted_token_ids, experts_ids,
                                      num_tokens_post_pad)


def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
                 token_expert_indicies: torch.Tensor,
                 gating_output: float) -> None:
    torch.ops._moe_C.topk_softmax(topk_weights, topk_ids,
                                  token_expert_indicies, gating_output)
330
331
332
333
334
335
336
337
338
339
340


def reshape_and_cache(
    key: torch.Tensor,
    value: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    kv_cache_dtype: str,
    kv_scale: float,
) -> None:
341
342
343
    torch.ops._C_cache_ops.reshape_and_cache(key, value, key_cache,
                                             value_cache, slot_mapping,
                                             kv_cache_dtype, kv_scale)
344
345


346
347
348
349
350
351
352
353
def reshape_and_cache_flash(
    key: torch.Tensor,
    value: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    kv_cache_dtype: str,
) -> None:
354
355
356
    torch.ops._C_cache_ops.reshape_and_cache_flash(key, value, key_cache,
                                                   value_cache, slot_mapping,
                                                   kv_cache_dtype)
357
358


359
360
def copy_blocks(key_caches: torch.Tensor, value_caches: torch.Tensor,
                block_mapping: torch.Tensor) -> None:
361
    torch.ops._C_cache_ops.copy_blocks(key_caches, value_caches, block_mapping)
362
363
364


def swap_blocks(src: torch.Tensor, dst: torch.Tensor,
365
                block_mapping: torch.Tensor) -> None:
366
    torch.ops._C_cache_ops.swap_blocks(src, dst, block_mapping)
367
368


369
370
371
372
def convert_fp8(output: torch.Tensor,
                input: torch.Tensor,
                scale: float = 1.0,
                kv_dtype: str = "fp8") -> None:
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
    torch.ops._C_cache_ops.convert_fp8(output, input, scale, kv_dtype)


def get_device_attribute(attribute: int, device: int) -> int:
    return torch.ops._C_cuda_utils.get_device_attribute(attribute, device)


def get_max_shared_memory_per_block_device_attribute(device: int) -> int:
    # ruff: noqa: E501
    return torch.ops._C_cuda_utils.get_max_shared_memory_per_block_device_attribute(
        device)


# custom ar
def init_custom_ar(meta: torch.Tensor, rank_data: torch.Tensor,
                   handles: List[str], offsets: List[int], rank: int,
                   full_nvlink: bool) -> int:
    return torch.ops._C_custom_ar.init_custom_ar(meta, rank_data, handles,
                                                 offsets, rank, full_nvlink)


def should_custom_ar(inp: torch.Tensor, max_size: int, world_size: int,
                     full_nvlink: bool) -> bool:
    return torch.ops._C_custom_ar.should_custom_ar(inp, max_size, world_size,
                                                   full_nvlink)


def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None:
    torch.ops._C_custom_ar.all_reduce_reg(fa, inp, out)

403

404
405
406
def all_reduce_unreg(fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor,
                     out: torch.Tensor) -> None:
    torch.ops._C_custom_ar.all_reduce_unreg(fa, inp, reg_buffer, out)
407

408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465

def dispose(fa: int) -> None:
    torch.ops._C_custom_ar.dispose(fa)


def meta_size() -> int:
    return torch.ops._C_custom_ar.meta_size()


def register_buffer(fa: int, t: torch.Tensor, handles: List[str],
                    offsets: List[int]) -> None:
    return torch.ops._C_custom_ar.register_buffer(fa, t, handles, offsets)


def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[str], List[int]]:
    return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta(fa)


def register_graph_buffers(fa: int, handles: List[str],
                           offsets: List[List[int]]) -> None:
    torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets)


# punica
def dispatch_bgmv(
    y: torch.Tensor,
    x: torch.Tensor,
    w_t_all: torch.Tensor,
    indicies: torch.Tensor,
    layer_idx: int,
    scale: float,
) -> None:
    torch.ops._punica_C.dispatch_bgmv(y, x, w_t_all, indicies, layer_idx,
                                      scale)


def dispatch_bgmv_low_level(
    y: torch.Tensor,
    x: torch.Tensor,
    w_t_all: torch.Tensor,
    indicies: torch.Tensor,
    layer_idx: int,
    scale: float,
    h_in: int,
    h_out: int,
    y_offset: int,
) -> None:
    torch.ops._punica_C.dispatch_bgmv_low_level(
        y,
        x,
        w_t_all,
        indicies,
        layer_idx,
        scale,
        h_in,
        h_out,
        y_offset,
    )