mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-11 19:55:01 +08:00
993 lines
39 KiB
Python
993 lines
39 KiB
Python
import contextlib
|
|
import functools
|
|
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
|
|
|
|
import torch
|
|
import torch.library
|
|
|
|
import vllm.envs as envs
|
|
from vllm._core_ext import ScalarType
|
|
from vllm.logger import init_logger
|
|
from vllm.platforms import current_platform
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
if not current_platform.is_tpu():
|
|
try:
|
|
import vllm._C
|
|
except ImportError as e:
|
|
logger.warning("Failed to import from vllm._C with %r", e)
|
|
|
|
if current_platform.is_rocm():
|
|
import vllm._rocm_C # noqa: F401
|
|
|
|
supports_moe_ops = False
|
|
with contextlib.suppress(ImportError):
|
|
import vllm._moe_C # noqa: F401
|
|
supports_moe_ops = True
|
|
|
|
if TYPE_CHECKING:
|
|
|
|
def register_fake(fn):
|
|
return lambda name: fn
|
|
else:
|
|
try:
|
|
from torch.library import register_fake
|
|
except ImportError:
|
|
from torch.library import impl_abstract as register_fake
|
|
|
|
|
|
def hint_on_error(fn):
|
|
|
|
@functools.wraps(fn)
|
|
def wrapper(*args, **kwargs):
|
|
try:
|
|
return fn(*args, **kwargs)
|
|
|
|
except NotImplementedError as e:
|
|
msg = (
|
|
"Error in calling custom op %s: %s\n"
|
|
"Not implemented or built, mostly likely because the current current device "
|
|
"does not support this kernel (less likely TORCH_CUDA_ARCH_LIST was set "
|
|
"incorrectly while building)")
|
|
logger.error(msg, fn.__name__, e)
|
|
raise NotImplementedError(msg % (fn.__name__, e)) from e
|
|
except AttributeError as e:
|
|
msg = (
|
|
"Error in calling custom op %s: %s\n"
|
|
"Possibly you have built or installed an obsolete version of vllm.\n"
|
|
"Please try a clean build and install of vllm,"
|
|
"or remove old built files such as vllm/*cpython*.so and build/ ."
|
|
)
|
|
logger.error(msg, fn.__name__, e)
|
|
raise e
|
|
|
|
return wrapper
|
|
|
|
|
|
# activation ops
|
|
def silu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
|
torch.ops._C.silu_and_mul(out, x)
|
|
|
|
|
|
def gelu_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
|
torch.ops._C.gelu_and_mul(out, x)
|
|
|
|
|
|
def gelu_tanh_and_mul(out: torch.Tensor, x: torch.Tensor) -> None:
|
|
torch.ops._C.gelu_tanh_and_mul(out, x)
|
|
|
|
|
|
def gelu_fast(out: torch.Tensor, x: torch.Tensor) -> None:
|
|
torch.ops._C.gelu_fast(out, x)
|
|
|
|
|
|
def gelu_new(out: torch.Tensor, x: torch.Tensor) -> None:
|
|
torch.ops._C.gelu_new(out, x)
|
|
|
|
|
|
def gelu_quick(out: torch.Tensor, x: torch.Tensor) -> None:
|
|
torch.ops._C.gelu_quick(out, x)
|
|
|
|
|
|
# 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,
|
|
seq_lens: torch.Tensor,
|
|
block_size: int,
|
|
max_seq_len: int,
|
|
alibi_slopes: Optional[torch.Tensor],
|
|
kv_cache_dtype: str,
|
|
k_scale: float,
|
|
v_scale: float,
|
|
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,
|
|
) -> None:
|
|
torch.ops._C.paged_attention_v1(
|
|
out, query, key_cache, value_cache, num_kv_heads, scale, block_tables,
|
|
seq_lens, block_size, max_seq_len, alibi_slopes, kv_cache_dtype,
|
|
k_scale, v_scale, tp_rank, blocksparse_local_blocks,
|
|
blocksparse_vert_stride, blocksparse_block_size,
|
|
blocksparse_head_sliding_step)
|
|
|
|
|
|
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,
|
|
seq_lens: torch.Tensor,
|
|
block_size: int,
|
|
max_seq_len: int,
|
|
alibi_slopes: Optional[torch.Tensor],
|
|
kv_cache_dtype: str,
|
|
k_scale: float,
|
|
v_scale: float,
|
|
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,
|
|
) -> None:
|
|
torch.ops._C.paged_attention_v2(
|
|
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, k_scale, v_scale, tp_rank,
|
|
blocksparse_local_blocks, blocksparse_vert_stride,
|
|
blocksparse_block_size, blocksparse_head_sliding_step)
|
|
|
|
|
|
def paged_attention_rocm(
|
|
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,
|
|
seq_lens: torch.Tensor,
|
|
block_size: int,
|
|
max_seq_len: int,
|
|
alibi_slopes: Optional[torch.Tensor],
|
|
kv_cache_dtype: str,
|
|
k_scale: float,
|
|
v_scale: float,
|
|
) -> None:
|
|
torch.ops._rocm_C.paged_attention(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, k_scale, v_scale)
|
|
|
|
|
|
# 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:
|
|
torch.ops._C.rotary_embedding(positions, query, key, head_size,
|
|
cos_sin_cache, is_neox)
|
|
|
|
|
|
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:
|
|
torch.ops._C.batched_rotary_embedding(positions, query, key, head_size,
|
|
cos_sin_cache, is_neox, rot_dim,
|
|
cos_sin_cache_offsets)
|
|
|
|
|
|
# layer norm ops
|
|
def rms_norm(out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor,
|
|
epsilon: float) -> None:
|
|
torch.ops._C.rms_norm(out, input, weight, epsilon)
|
|
|
|
|
|
def fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor,
|
|
weight: torch.Tensor, epsilon: float) -> None:
|
|
torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon)
|
|
|
|
|
|
def advance_step_flashattn(num_seqs: int, num_queries: int, block_size: int,
|
|
input_tokens: torch.Tensor,
|
|
sampled_token_ids: torch.Tensor,
|
|
input_positions: torch.Tensor,
|
|
seq_lens: torch.Tensor, slot_mapping: torch.Tensor,
|
|
block_tables: torch.Tensor) -> None:
|
|
"""Advance a step on GPU for existing inputs for a multi-step runner"""
|
|
return torch.ops._C.advance_step_flashattn(num_seqs, num_queries,
|
|
block_size, input_tokens,
|
|
sampled_token_ids,
|
|
input_positions, seq_lens,
|
|
slot_mapping, block_tables)
|
|
|
|
|
|
def advance_step_flashinfer(num_seqs: int, num_queries: int, block_size: int,
|
|
input_tokens: torch.Tensor,
|
|
sampled_token_ids: torch.Tensor,
|
|
input_positions: torch.Tensor,
|
|
seq_lens: torch.Tensor, slot_mapping: torch.Tensor,
|
|
block_tables: torch.Tensor,
|
|
paged_kv_indices: torch.Tensor,
|
|
paged_kv_indptr: torch.Tensor,
|
|
paged_kv_last_page_len: torch.Tensor,
|
|
block_table_bound: torch.Tensor) -> None:
|
|
|
|
return torch.ops._C.advance_step_flashinfer(
|
|
num_seqs, num_queries, block_size, input_tokens, sampled_token_ids,
|
|
input_positions, seq_lens, slot_mapping, block_tables,
|
|
paged_kv_indices, paged_kv_indptr, paged_kv_last_page_len,
|
|
block_table_bound)
|
|
|
|
|
|
# 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:
|
|
if envs.VLLM_USE_TRITON_AWQ:
|
|
from vllm.model_executor.layers.quantization.awq_triton import (
|
|
awq_dequantize_triton)
|
|
return awq_dequantize_triton(qweight, scales, zeros)
|
|
return torch.ops._C.awq_dequantize(qweight, scales, zeros, split_k_iters,
|
|
thx, thy)
|
|
|
|
|
|
def awq_gemm(input: torch.Tensor, qweight: torch.Tensor, qzeros: torch.Tensor,
|
|
scales: torch.Tensor, split_k_iters: int) -> torch.Tensor:
|
|
if envs.VLLM_USE_TRITON_AWQ:
|
|
from vllm.model_executor.layers.quantization.awq_triton import (
|
|
awq_gemm_triton)
|
|
return awq_gemm_triton(input, qweight, qzeros, scales, split_k_iters)
|
|
return torch.ops._C.awq_gemm(input, qweight, qzeros, scales, split_k_iters)
|
|
|
|
|
|
# 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:
|
|
return torch.ops._C.gptq_gemm(a, b_q_weight, b_gptq_qzeros, b_gptq_scales,
|
|
b_g_idx, use_exllama, bit)
|
|
|
|
|
|
if hasattr(torch.ops._C, "gptq_gemm"):
|
|
|
|
@register_fake("_C::gptq_gemm")
|
|
def _gptq_gemm_fake(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:
|
|
return torch.empty((a.size(0), b_q_weight.size(1)),
|
|
dtype=a.dtype,
|
|
device=a.device)
|
|
|
|
|
|
def gptq_shuffle(q_weight: torch.Tensor, q_perm: torch.Tensor,
|
|
bit: int) -> None:
|
|
torch.ops._C.gptq_shuffle(q_weight, q_perm, bit)
|
|
|
|
|
|
# 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:
|
|
return torch.ops._C.marlin_gemm(a, b_q_weight, b_scales, workspace, size_m,
|
|
size_n, size_k)
|
|
|
|
|
|
# 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, b_q_type: ScalarType,
|
|
size_m: int, size_n: int, size_k: int) -> torch.Tensor:
|
|
return torch.ops._C.gptq_marlin_24_gemm(a, b_q_weight, b_meta, b_scales,
|
|
workspace, b_q_type, size_m,
|
|
size_n, size_k)
|
|
|
|
|
|
if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
|
|
|
|
@register_fake("_C::gptq_marlin_24_gemm")
|
|
def _gptq_marlin_24_gemm_fake(a: torch.Tensor, b_q_weight: torch.Tensor,
|
|
b_meta: torch.Tensor, b_scales: torch.Tensor,
|
|
workspace: torch.Tensor,
|
|
b_q_type: ScalarType, size_m: int,
|
|
size_n: int, size_k: int) -> torch.Tensor:
|
|
return torch.empty((size_m, size_n), device=a.device, dtype=a.dtype)
|
|
|
|
@register_fake("_C::gptq_marlin_gemm")
|
|
def _gptq_marlin_gemm_fake(a: torch.Tensor,
|
|
b_q_weight: torch.Tensor,
|
|
b_scales: torch.Tensor,
|
|
b_zeros: torch.Tensor,
|
|
g_idx: torch.Tensor,
|
|
perm: torch.Tensor,
|
|
workspace: torch.Tensor,
|
|
b_q_type: ScalarType,
|
|
size_m: int,
|
|
size_n: int,
|
|
size_k: int,
|
|
is_k_full: bool,
|
|
has_zp: bool = False,
|
|
use_fp32_reduce: bool = False) -> torch.Tensor:
|
|
return torch.empty((size_m, size_n), device=a.device, dtype=a.dtype)
|
|
|
|
@register_fake("_C::ggml_dequantize")
|
|
def _ggml_dequantize_fake(W: torch.Tensor, quant_type: int, m: int,
|
|
n: int) -> torch.Tensor:
|
|
return torch.empty((m, n), dtype=torch.float16, device=W.device)
|
|
|
|
@register_fake("_C::ggml_mul_mat_vec_a8")
|
|
def _ggml_mul_mat_vec_a8_fake(
|
|
W: torch.Tensor,
|
|
X: torch.Tensor,
|
|
quant_type: int,
|
|
row: int,
|
|
) -> torch.Tensor:
|
|
return torch.empty((1, row), dtype=torch.float16, device=W.device)
|
|
|
|
@register_fake("_C::ggml_mul_mat_a8")
|
|
def _ggml_mul_mat_a8_fake(
|
|
W: torch.Tensor,
|
|
X: torch.Tensor,
|
|
quant_type: int,
|
|
row: int,
|
|
) -> torch.Tensor:
|
|
batch = X.size(0)
|
|
return torch.empty((batch, row), dtype=torch.float16, device=W.device)
|
|
|
|
@register_fake("_C::marlin_qqq_gemm")
|
|
def _marlin_qqq_gemm_fake(a: torch.Tensor, b_q_weight: torch.Tensor,
|
|
s_tok: torch.Tensor, s_ch: torch.Tensor,
|
|
s_group: torch.Tensor, workspace: torch.Tensor,
|
|
size_m: int, size_n: int,
|
|
size_k: int) -> torch.Tensor:
|
|
return torch.empty((size_m, size_n),
|
|
dtype=torch.float16,
|
|
device=a.device)
|
|
|
|
@register_fake("_C::marlin_gemm")
|
|
def _marlin_gemm_fake(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:
|
|
return torch.empty((size_m, size_n),
|
|
dtype=torch.float16,
|
|
device=a.device)
|
|
|
|
@register_fake("_C::awq_dequantize")
|
|
def _awq_dequantize_fake(qweight: torch.Tensor, scales: torch.Tensor,
|
|
zeros: torch.Tensor, split_k_iters: int, thx: int,
|
|
thy: int) -> torch.Tensor:
|
|
in_c = qweight.size(0)
|
|
qout_c = qweight.size(1)
|
|
out_c = qout_c * 8
|
|
return torch.empty((in_c, out_c),
|
|
dtype=scales.dtype,
|
|
device=scales.device)
|
|
|
|
@register_fake("_C::awq_gemm")
|
|
def _awq_gemm_fake(input: torch.Tensor, qweight: torch.Tensor,
|
|
qzeros: torch.Tensor, scales: torch.Tensor,
|
|
split_k_iters: int) -> torch.Tensor:
|
|
num_in_feats = input.size(0)
|
|
return torch.empty((split_k_iters, num_in_feats, qweight.size(1) * 8),
|
|
dtype=input.dtype,
|
|
device=input.device).sum(0)
|
|
|
|
@register_fake("_C::aqlm_gemm")
|
|
def _aqlm_gemm_fake(input: torch.Tensor, codes: torch.Tensor,
|
|
codebooks: torch.Tensor, scales: torch.Tensor,
|
|
codebook_partition_sizes: List[int],
|
|
bias: Optional[torch.Tensor]) -> torch.Tensor:
|
|
out_features = codes.size(0) * codebooks.size(2)
|
|
flat_input = input.reshape((-1, input.size(-1)))
|
|
flat_output = torch.empty((flat_input.size(0), out_features),
|
|
dtype=input.dtype,
|
|
device=input.device)
|
|
|
|
output_sizes = list(input.shape)
|
|
output_sizes.pop()
|
|
output_sizes.append(-1)
|
|
return flat_output.reshape(tuple(output_sizes))
|
|
|
|
@register_fake("_C::aqlm_dequant")
|
|
def _aqlm_dequant_fake(
|
|
codes: torch.Tensor, codebooks: torch.Tensor,
|
|
codebook_partition_sizes: List[int]) -> torch.Tensor:
|
|
in_features = codes.size(1) * 8
|
|
out_features = codes.size(0)
|
|
return torch.empty((out_features, in_features),
|
|
dtype=codebooks.dtype,
|
|
device=codebooks.device)
|
|
|
|
@register_fake("_C::fp8_marlin_gemm")
|
|
def _fp8_marlin_gemm_fake(a: torch.Tensor, b_q_weight: torch.Tensor,
|
|
b_scales: torch.Tensor, workspace: torch.Tensor,
|
|
num_bits: int, size_m: int, size_n: int,
|
|
size_k: int) -> torch.Tensor:
|
|
return torch.empty((size_m, size_n), dtype=a.dtype, device=a.device)
|
|
|
|
@register_fake("_C::machete_gemm")
|
|
def machete_gemm_fake(
|
|
a: torch.Tensor,
|
|
# Should be the tensor returned by machete_prepack_B
|
|
b_q: torch.Tensor,
|
|
b_type: ScalarType,
|
|
b_scales: Optional[torch.Tensor] = None,
|
|
b_zeros: Optional[torch.Tensor] = None,
|
|
b_group_size: Optional[int] = None,
|
|
c: Optional[torch.Tensor] = None,
|
|
alpha: Optional[float] = None,
|
|
beta: Optional[float] = None,
|
|
schedule: Optional[str] = None,
|
|
) -> torch.Tensor:
|
|
m = a.size(0)
|
|
n = b_q.size(1)
|
|
return torch.empty((m, n), device=a.device, dtype=a.dtype)
|
|
|
|
@register_fake("_C::machete_prepack_B")
|
|
def machete_prepack_B_fake(b_q_weight: torch.Tensor,
|
|
b_type: ScalarType) -> torch.Tensor:
|
|
return torch.empty_like(b_q_weight,
|
|
memory_format=torch.contiguous_format)
|
|
|
|
@register_fake("_C::causal_conv1d_fwd")
|
|
def causal_conv1d_fwd_fake(x: torch.Tensor, weight: torch.Tensor,
|
|
bias_: Optional[torch.Tensor],
|
|
conv_states: Optional[torch.Tensor],
|
|
cu_seq_len: Optional[torch.Tensor],
|
|
cache_indices: Optional[torch.Tensor],
|
|
has_initial_state: Optional[torch.Tensor],
|
|
silu_activation: bool, pad_slot_id: int):
|
|
return None
|
|
|
|
@register_fake("_C::causal_conv1d_update")
|
|
def causal_conv1d_update_fake(x: torch.Tensor, conv_state: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
bias_: Optional[torch.Tensor],
|
|
silu_activation: bool,
|
|
cache_seqlens: Optional[torch.Tensor],
|
|
conv_state_indices: Optional[torch.Tensor],
|
|
pad_slot_id: int) -> None:
|
|
return None
|
|
|
|
@register_fake("_C::selective_scan_fwd")
|
|
def selective_scan_fwd_fake(u: torch.Tensor, delta: torch.Tensor,
|
|
A: torch.Tensor, B: torch.Tensor,
|
|
C: torch.Tensor, D_: Optional[torch.Tensor],
|
|
z_: Optional[torch.Tensor],
|
|
delta_bias_: Optional[torch.Tensor],
|
|
delta_softplus: bool,
|
|
cu_seq_len: Optional[torch.Tensor],
|
|
cache_indices: Optional[torch.Tensor],
|
|
has_initial_state: Optional[torch.Tensor],
|
|
ssm_states: Optional[torch.Tensor],
|
|
pad_slot_id: int) -> None:
|
|
return None
|
|
|
|
|
|
# cutlass
|
|
def cutlass_scaled_mm_supports_fp8(cuda_device_capability: int) -> bool:
|
|
return torch.ops._C.cutlass_scaled_mm_supports_fp8(cuda_device_capability)
|
|
|
|
|
|
def cutlass_scaled_mm(a: torch.Tensor,
|
|
b: torch.Tensor,
|
|
scale_a: torch.Tensor,
|
|
scale_b: torch.Tensor,
|
|
out_dtype: torch.dtype,
|
|
bias: Optional[torch.Tensor] = None) -> 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)
|
|
assert bias is None or bias.shape[0] == b.shape[
|
|
1] and bias.dtype == out_dtype
|
|
|
|
m = a.shape[0]
|
|
n = b.shape[1]
|
|
out = torch.empty((m, n), dtype=out_dtype, device=a.device)
|
|
|
|
torch.ops._C.cutlass_scaled_mm(out, a, b, scale_a, scale_b, bias)
|
|
|
|
return out
|
|
|
|
|
|
def cutlass_scaled_mm_azp(a: torch.Tensor,
|
|
b: torch.Tensor,
|
|
scale_a: torch.Tensor,
|
|
scale_b: torch.Tensor,
|
|
out_dtype: torch.dtype,
|
|
azp_adj: torch.Tensor,
|
|
azp: Optional[torch.Tensor] = None,
|
|
bias: Optional[torch.Tensor] = None) -> 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)
|
|
assert bias is None or bias.numel(
|
|
) == b.shape[1] and bias.dtype == out_dtype
|
|
|
|
m = a.shape[0]
|
|
n = b.shape[1]
|
|
out = torch.empty((m, n), dtype=out_dtype, device=a.device)
|
|
|
|
torch.ops._C.cutlass_scaled_mm_azp(out, a, b, scale_a, scale_b, azp_adj,
|
|
azp, bias)
|
|
return out
|
|
|
|
|
|
# aqlm
|
|
def aqlm_gemm(input: torch.Tensor, codes: torch.Tensor,
|
|
codebooks: torch.Tensor, scales: torch.Tensor,
|
|
codebook_partition_sizes: List[int],
|
|
bias: Optional[torch.Tensor]) -> torch.Tensor:
|
|
return torch.ops._C.aqlm_gemm(input, codes, codebooks, scales,
|
|
codebook_partition_sizes, bias)
|
|
|
|
|
|
def aqlm_dequant(codes: torch.Tensor, codebooks: torch.Tensor,
|
|
codebook_partition_sizes: List[int]) -> torch.Tensor:
|
|
return torch.ops._C.aqlm_dequant(codes, codebooks,
|
|
codebook_partition_sizes)
|
|
|
|
|
|
# gptq_marlin
|
|
def gptq_marlin_repack(b_q_weight: torch.Tensor, perm: torch.Tensor,
|
|
size_k: int, size_n: int,
|
|
num_bits: int) -> torch.Tensor:
|
|
return torch.ops._C.gptq_marlin_repack(b_q_weight, perm, size_k, size_n,
|
|
num_bits)
|
|
|
|
|
|
# gptq_marlin
|
|
def awq_marlin_repack(b_q_weight: torch.Tensor, size_k: int, size_n: int,
|
|
num_bits: int) -> torch.Tensor:
|
|
return torch.ops._C.awq_marlin_repack(b_q_weight, size_k, size_n, num_bits)
|
|
|
|
|
|
def gptq_marlin_moe_repack(b_q_weight: torch.Tensor, perm: torch.Tensor,
|
|
size_k: int, size_n: int,
|
|
num_bits: int) -> torch.Tensor:
|
|
num_experts = b_q_weight.shape[0]
|
|
assert size_k % 16 == 0
|
|
output = torch.empty((num_experts, size_k // 16, size_n * (num_bits // 2)),
|
|
device=b_q_weight.device,
|
|
dtype=b_q_weight.dtype)
|
|
for e in range(num_experts):
|
|
output[e] = torch.ops._C.gptq_marlin_repack(b_q_weight[e], perm[e],
|
|
size_k, size_n, num_bits)
|
|
return output
|
|
|
|
|
|
def awq_marlin_moe_repack(b_q_weight: torch.Tensor, perm: torch.Tensor,
|
|
size_k: int, size_n: int,
|
|
num_bits: int) -> torch.Tensor:
|
|
num_experts = b_q_weight.shape[0]
|
|
assert size_k % 16 == 0
|
|
output = torch.empty((num_experts, size_k // 16, size_n * (num_bits // 2)),
|
|
device=b_q_weight.device,
|
|
dtype=b_q_weight.dtype)
|
|
for e in range(num_experts):
|
|
output[e] = torch.ops._C.awq_marlin_repack(b_q_weight[e], size_k,
|
|
size_n, num_bits)
|
|
return output
|
|
|
|
|
|
def gptq_marlin_gemm(a: torch.Tensor,
|
|
b_q_weight: torch.Tensor,
|
|
b_scales: torch.Tensor,
|
|
b_zeros: torch.Tensor,
|
|
g_idx: torch.Tensor,
|
|
perm: torch.Tensor,
|
|
workspace: torch.Tensor,
|
|
b_q_type: ScalarType,
|
|
size_m: int,
|
|
size_n: int,
|
|
size_k: int,
|
|
is_k_full: bool,
|
|
has_zp: bool = False,
|
|
use_fp32_reduce: bool = False) -> torch.Tensor:
|
|
return torch.ops._C.gptq_marlin_gemm(a, b_q_weight, b_scales, b_zeros,
|
|
g_idx, perm, workspace, b_q_type,
|
|
size_m, size_n, size_k, is_k_full,
|
|
has_zp, use_fp32_reduce)
|
|
|
|
|
|
# fp8 marlin
|
|
def fp8_marlin_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
|
|
b_scales: torch.Tensor, workspace: torch.Tensor,
|
|
num_bits: int, size_m: int, size_n: int,
|
|
size_k: int) -> torch.Tensor:
|
|
return torch.ops._C.fp8_marlin_gemm(a, b_q_weight, b_scales, workspace,
|
|
num_bits, size_m, size_n, size_k)
|
|
|
|
|
|
# machete
|
|
def machete_supported_schedules(b_type: ScalarType) -> List[str]:
|
|
return torch.ops._C.machete_supported_schedules(b_type)
|
|
|
|
|
|
def machete_gemm(
|
|
a: torch.Tensor,
|
|
b_q: torch.Tensor, # Should be the tensor returned by machete_prepack_B
|
|
b_type: ScalarType,
|
|
b_scales: Optional[torch.Tensor] = None,
|
|
b_zeros: Optional[torch.Tensor] = None,
|
|
b_group_size: Optional[int] = None,
|
|
c: Optional[torch.Tensor] = None,
|
|
alpha: Optional[float] = None,
|
|
beta: Optional[float] = None,
|
|
schedule: Optional[str] = None,
|
|
) -> torch.Tensor:
|
|
return torch.ops._C.machete_gemm(a, b_q, b_type, b_scales, b_zeros,
|
|
b_group_size, c, alpha, beta, schedule)
|
|
|
|
|
|
def machete_prepack_B(b_q_weight: torch.Tensor,
|
|
b_type: ScalarType) -> torch.Tensor:
|
|
return torch.ops._C.machete_prepack_B(b_q_weight, b_type)
|
|
|
|
|
|
if hasattr(torch.ops._C, "permute_cols"):
|
|
|
|
@register_fake("_C::permute_cols")
|
|
def _permute_cols_fake(a: torch.Tensor,
|
|
perm: torch.Tensor) -> torch.Tensor:
|
|
return torch.empty_like(a)
|
|
|
|
|
|
def permute_cols(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:
|
|
return torch.ops._C.permute_cols(a, perm)
|
|
|
|
|
|
# fp8
|
|
def scaled_fp8_quant(
|
|
input: torch.Tensor,
|
|
scale: Optional[torch.Tensor] = None,
|
|
num_token_padding: Optional[int] = None,
|
|
scale_ub: Optional[torch.Tensor] = None,
|
|
use_per_token_if_dynamic: bool = False,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
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 tensors 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
|
|
scale_ub: Optional upper bound for scaling factor in dynamic
|
|
per token case
|
|
num_token_padding: If specified, pad the first dimension
|
|
of the output to at least this value.
|
|
use_per_token_if_dynamic: Whether to do per_tensor or per_token
|
|
in the dynamic quantization case.
|
|
|
|
Returns:
|
|
Tuple[torch.Tensor, torch.Tensor]: The output tensor in FP8 and
|
|
scaling factor.
|
|
"""
|
|
# This code assumes batch_dim and num_tokens are flattened
|
|
assert (input.ndim == 2)
|
|
shape: Union[Tuple[int, int], torch.Size] = input.shape
|
|
# For rocm, the output fp8 dtype is torch.float_e3m3fnuz
|
|
out_dtype: torch.dtype = torch.float8_e4m3fnuz if vllm.utils.is_hip() \
|
|
else torch.float8_e4m3fn
|
|
if num_token_padding:
|
|
shape = (max(num_token_padding, input.shape[0]), shape[1])
|
|
output = torch.empty(shape, device=input.device, dtype=out_dtype)
|
|
|
|
if scale is None:
|
|
if use_per_token_if_dynamic:
|
|
scale = torch.empty((shape[0], 1),
|
|
device=input.device,
|
|
dtype=torch.float32)
|
|
torch.ops._C.dynamic_per_token_scaled_fp8_quant(
|
|
output, input, scale, scale_ub)
|
|
else:
|
|
scale = torch.zeros(1, device=input.device, dtype=torch.float32)
|
|
torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)
|
|
else:
|
|
# num_token_padding not implemented for this case
|
|
assert (scale.numel() == 1 or num_token_padding is None)
|
|
torch.ops._C.static_scaled_fp8_quant(output, input, scale)
|
|
|
|
return output, scale
|
|
|
|
|
|
# int8
|
|
def scaled_int8_quant(
|
|
input: torch.Tensor,
|
|
scale: Optional[torch.Tensor] = None,
|
|
azp: Optional[torch.Tensor] = None,
|
|
symmetric: bool = True
|
|
) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
|
|
"""
|
|
Quantize the input tensor to int8 and return the quantized tensor and scale, and maybe azp.
|
|
|
|
Args:
|
|
input: The input tensor to be quantized to int8.
|
|
scale: Optional scaling factor for the int8 quantization.
|
|
When not provided, we invoke dynamic-per-token quantization.
|
|
azp: Optional zero-point for the int8 quantization.
|
|
Must be provided for asymmetric quantization if `scale` is provided.
|
|
symmetric: Whether to use symmetric quantization (scale only, azp ignored).
|
|
|
|
Returns:
|
|
Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]] : Output int8 tensor, scales, and optionally azp.
|
|
"""
|
|
output = torch.empty_like(input, dtype=torch.int8)
|
|
if scale is not None:
|
|
# static-per-tensor quantization.
|
|
assert symmetric == (
|
|
azp is
|
|
None), "azp must only be provided for asymmetric quantization."
|
|
torch.ops._C.static_scaled_int8_quant(output, input, scale, azp)
|
|
return output, scale, None
|
|
|
|
# dynamic-per-token quantization.
|
|
input_scales = torch.empty((input.numel() // input.shape[-1], 1),
|
|
device=input.device,
|
|
dtype=torch.float32)
|
|
input_azp = None if symmetric else torch.empty_like(input_scales,
|
|
dtype=torch.int32)
|
|
torch.ops._C.dynamic_scaled_int8_quant(output, input, input_scales,
|
|
input_azp)
|
|
return output, input_scales, input_azp
|
|
|
|
|
|
# qqq ops
|
|
def marlin_qqq_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
|
|
s_tok: torch.Tensor, s_ch: torch.Tensor,
|
|
s_group: torch.Tensor, workspace: torch.Tensor,
|
|
size_m: int, size_n: int, size_k: int) -> torch.Tensor:
|
|
return torch.ops._C.marlin_qqq_gemm(a, b_q_weight, s_tok, s_ch, s_group,
|
|
workspace, size_m, size_n, size_k)
|
|
|
|
|
|
# gguf
|
|
def ggml_dequantize(W: torch.Tensor, quant_type: int, m: int,
|
|
n: int) -> torch.Tensor:
|
|
return torch.ops._C.ggml_dequantize(W, quant_type, m, n)
|
|
|
|
|
|
def ggml_mul_mat_vec_a8(
|
|
W: torch.Tensor,
|
|
X: torch.Tensor,
|
|
quant_type: int,
|
|
row: int,
|
|
) -> torch.Tensor:
|
|
return torch.ops._C.ggml_mul_mat_vec_a8(W, X, quant_type, row)
|
|
|
|
|
|
def ggml_mul_mat_a8(
|
|
W: torch.Tensor,
|
|
X: torch.Tensor,
|
|
quant_type: int,
|
|
row: int,
|
|
) -> torch.Tensor:
|
|
return torch.ops._C.ggml_mul_mat_a8(W, X, quant_type, row)
|
|
|
|
|
|
# mamba
|
|
def causal_conv1d_fwd(x: torch.Tensor, weight: torch.Tensor,
|
|
bias_: Optional[torch.Tensor],
|
|
conv_states: Optional[torch.Tensor],
|
|
query_start_loc: Optional[torch.Tensor],
|
|
cache_indices: Optional[torch.Tensor],
|
|
has_initial_state: Optional[torch.Tensor],
|
|
silu_activation: bool, pad_slot_id: int):
|
|
torch.ops._C.causal_conv1d_fwd(x, weight, bias_, conv_states,
|
|
query_start_loc, cache_indices,
|
|
has_initial_state, silu_activation,
|
|
pad_slot_id)
|
|
|
|
|
|
def causal_conv1d_update(x: torch.Tensor, conv_state: torch.Tensor,
|
|
weight: torch.Tensor, bias_: Optional[torch.Tensor],
|
|
silu_activation: bool,
|
|
cache_seqlens: Optional[torch.Tensor],
|
|
conv_state_indices: Optional[torch.Tensor],
|
|
pad_slot_id: int):
|
|
torch.ops._C.causal_conv1d_update(x, conv_state, weight, bias_,
|
|
silu_activation, cache_seqlens,
|
|
conv_state_indices, pad_slot_id)
|
|
|
|
|
|
def selective_scan_fwd(u: torch.Tensor, delta: torch.Tensor, A: torch.Tensor,
|
|
B: torch.Tensor, C: torch.Tensor,
|
|
D_: Optional[torch.Tensor], z_: Optional[torch.Tensor],
|
|
delta_bias_: Optional[torch.Tensor],
|
|
delta_softplus: bool,
|
|
query_start_loc: Optional[torch.Tensor],
|
|
cache_indices: Optional[torch.Tensor],
|
|
has_initial_state: Optional[torch.Tensor],
|
|
ssm_states: torch.Tensor, pad_slot_id: int):
|
|
torch.ops._C.selective_scan_fwd(u, delta, A, B, C, D_, z_, delta_bias_,
|
|
delta_softplus, query_start_loc,
|
|
cache_indices, has_initial_state,
|
|
ssm_states, pad_slot_id)
|
|
|
|
|
|
# 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:
|
|
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)
|
|
|
|
|
|
if supports_moe_ops and hasattr(torch.ops._moe_C, "marlin_gemm_moe"):
|
|
|
|
@register_fake("_moe_C::marlin_gemm_moe")
|
|
def marlin_gemm_moe_fake(a: torch.Tensor, b_q_weights: torch.Tensor,
|
|
sorted_ids: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor, b_scales: torch.Tensor,
|
|
b_zero_points: torch.Tensor, g_idx: torch.Tensor,
|
|
perm: torch.Tensor, workspace: torch.Tensor,
|
|
b_q_type: ScalarType, size_m: int, size_n: int,
|
|
size_k: int, is_k_full: bool, num_experts: int,
|
|
topk: int, moe_block_size: int,
|
|
replicate_input: bool,
|
|
apply_weights: bool) -> torch.Tensor:
|
|
return torch.empty((size_m, topk, size_n),
|
|
dtype=a.dtype,
|
|
device=a.device)
|
|
|
|
|
|
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,
|
|
k_scale: float,
|
|
v_scale: float,
|
|
) -> None:
|
|
torch.ops._C_cache_ops.reshape_and_cache(key, value, key_cache,
|
|
value_cache, slot_mapping,
|
|
kv_cache_dtype, k_scale, v_scale)
|
|
|
|
|
|
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,
|
|
k_scale: float,
|
|
v_scale: float,
|
|
) -> None:
|
|
torch.ops._C_cache_ops.reshape_and_cache_flash(key, value, key_cache,
|
|
value_cache, slot_mapping,
|
|
kv_cache_dtype, k_scale,
|
|
v_scale)
|
|
|
|
|
|
def copy_blocks(key_caches: List[torch.Tensor],
|
|
value_caches: List[torch.Tensor],
|
|
block_mapping: torch.Tensor) -> None:
|
|
torch.ops._C_cache_ops.copy_blocks(key_caches, value_caches, block_mapping)
|
|
|
|
|
|
def swap_blocks(src: torch.Tensor, dst: torch.Tensor,
|
|
block_mapping: torch.Tensor) -> None:
|
|
torch.ops._C_cache_ops.swap_blocks(src, dst, block_mapping)
|
|
|
|
|
|
def convert_fp8(output: torch.Tensor,
|
|
input: torch.Tensor,
|
|
scale: float = 1.0,
|
|
kv_dtype: str = "fp8") -> None:
|
|
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 all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None:
|
|
torch.ops._C_custom_ar.all_reduce_reg(fa, inp, out)
|
|
|
|
|
|
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)
|
|
|
|
|
|
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)
|
|
|
|
|
|
# temporary fix for https://github.com/vllm-project/vllm/issues/5456
|
|
# TODO: remove this in v0.6.0
|
|
names_and_values = globals()
|
|
names_and_values_to_update = {}
|
|
# prepare variables to avoid dict size change during iteration
|
|
k, v, arg = None, None, None
|
|
fn_type = type(lambda x: x)
|
|
for k, v in names_and_values.items():
|
|
# find functions that are defined in this file and have torch.Tensor
|
|
# in their annotations. `arg == "torch.Tensor"` is used to handle
|
|
# the case when users use `import __annotations__` to turn type
|
|
# hints into strings.
|
|
if isinstance(v, fn_type) \
|
|
and v.__code__.co_filename == __file__ \
|
|
and any(arg is torch.Tensor or arg == "torch.Tensor"
|
|
for arg in v.__annotations__.values()):
|
|
names_and_values_to_update[k] = hint_on_error(v)
|
|
|
|
names_and_values.update(names_and_values_to_update)
|
|
del names_and_values_to_update, names_and_values, v, k, fn_type
|