feat(CRITICAL): 从 GitHub 扫描搬运 ixformer SDK + xllm 完整 GDN/MoE 代码
来源:
1. Chranos/ixformer (GitHub) → ixformer_sdk/ (230 files, 70K lines)
- inference/functions/vllm.py: vllm_moe_topk_softmax 完整实现 (2033 lines)
- inference/functions/moe.py: MoE ops 完整实现 (1380 lines)
- contrib/vllm_flash_attn/: FA2 Python 接口 (1018 lines)
- contrib/tgi/fused_moe.py: TGI fused MoE (429 lines)
- csrc/include/ixformer/: C++ kernel headers + cmake
2. Deep-Spark/xllm (GitHub) → upstream_ref/xllm_latest/ (+15 files)
- npu_torch/qwen3_5_decoder_layer_impl.cpp/.h
- npu_torch/qwen3_5_gated_delta_net.cpp/.h
- npu_torch/qwen3_next_*.cpp/.h (6 files)
- npu_torch/attention.cpp/.h + fused_moe.cpp/.h + CMakeLists.txt
- models/llm/qwen3_5.h + qwen3_5_mtp.h + qwen3_next.h
- models/vlm/qwen3_5.h
调用链完整性:
ixformer_sdk/inference/functions/vllm.py
→ ops.infer.moe_topk_softmax() (C++ 层)
→ 这就是 base 镜像 libixformer.so 里的实现
upstream_ref/xllm_latest/core/layers/ilu/fused_moe.cpp
→ ixformer::infer::topk_softmax() (直接 C++ 调用)
→ ixformer::infer::group_gemm() → 完整 7-step MoE pipeline
This commit is contained in:
project6-dev
44
ixformer_sdk/inference/functions/__init__.py
Normal file
44
ixformer_sdk/inference/functions/__init__.py
Normal file
@@ -0,0 +1,44 @@
|
||||
from .act_and_mul import *
|
||||
from .act_bias_mm import *
|
||||
from .add import *
|
||||
from .bert import *
|
||||
from .bnb_dequant import *
|
||||
from .bnb_double_quant import *
|
||||
from .bnb_mm_dequant import *
|
||||
from .bnb_qgemm import *
|
||||
from .bnb_quant import *
|
||||
from .bnb_rowcol_absmax import *
|
||||
from .conv2d import *
|
||||
from .cross_entropy_loss import *
|
||||
from .flash_attn import *
|
||||
from .flash_attn_lib import *
|
||||
from .fused_rope import *
|
||||
from .gemv import *
|
||||
from .groupnorm import *
|
||||
from .i8w8o32 import *
|
||||
from .layernorm import *
|
||||
from .lightllm import *
|
||||
from .linalg import *
|
||||
from .linear import *
|
||||
from .lmdeploy import *
|
||||
from .marlin import *
|
||||
from .matmul import *
|
||||
from .mla_fused import *
|
||||
from .mm import *
|
||||
from .moe import *
|
||||
from .overlap_comm import *
|
||||
from .paged_attention import *
|
||||
from .quantized_linear import *
|
||||
from .residual_bias import *
|
||||
from .rms_norm import *
|
||||
from .scaled_dot_product_attention import *
|
||||
from .smoothquant import *
|
||||
from .softmax import *
|
||||
from .store_kv_cache import *
|
||||
from .t5 import *
|
||||
from .tgi import *
|
||||
from .vllm import *
|
||||
from .w8a8 import *
|
||||
from .w8a16 import *
|
||||
from .wi4a16 import *
|
||||
from .wui4a16 import *
|
||||
88
ixformer_sdk/inference/functions/act_and_mul.py
Normal file
88
ixformer_sdk/inference/functions/act_and_mul.py
Normal file
@@ -0,0 +1,88 @@
|
||||
from typing import List, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
import torch.nn.functional as NNF
|
||||
|
||||
__all__ = ["ref_silu_and_mul", "ref_gelu_and_mul", "ref_gelu_tanh_and_mul",
|
||||
"silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul"]
|
||||
|
||||
|
||||
def ref_silu_and_mul(input: "torch.Tensor") -> torch.Tensor:
|
||||
x1, x2 = input.chunk(chunks=2, dim=-1)
|
||||
res = NNF.silu(x1) * x2
|
||||
return res
|
||||
|
||||
|
||||
def ref_gelu_and_mul(input: "torch.Tensor", gate_first=True) -> torch.Tensor:
|
||||
x1, x2 = input.chunk(chunks=2, dim=-1)
|
||||
if gate_first:
|
||||
res = NNF.gelu(x1) * x2
|
||||
else:
|
||||
res = NNF.gelu(x2) * x1
|
||||
return res
|
||||
|
||||
|
||||
def ref_gelu_tanh_and_mul(input: "torch.Tensor") -> torch.Tensor:
|
||||
x1, x2 = input.chunk(chunks=2, dim=-1)
|
||||
res = NNF.gelu(x1) * x2
|
||||
return res
|
||||
|
||||
|
||||
def silu_and_mul(input: torch.Tensor, output: torch.Tensor = None):
|
||||
|
||||
"""
|
||||
Args:
|
||||
input: (..., 2*hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
|
||||
Returns:
|
||||
output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
"""
|
||||
if output is None:
|
||||
output_shape = list(input.shape)
|
||||
output_shape[-1] = output_shape[-1] // 2
|
||||
output = input.new_empty(output_shape)
|
||||
|
||||
ops.infer.silu_and_mul(input, output)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def gelu_and_mul(input: "torch.Tensor", output: torch.Tensor = None, gate_first=True):
|
||||
|
||||
"""
|
||||
Args:
|
||||
input: (..., 2*hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
gate_first: bool
|
||||
Returns:
|
||||
output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
"""
|
||||
if output is None:
|
||||
output_shape = list(input.shape)
|
||||
output_shape[-1] = output_shape[-1] // 2
|
||||
output = input.new_empty(output_shape)
|
||||
|
||||
ops.infer.gelu_and_mul(input, output, gate_first)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def gelu_tanh_and_mul(input: torch.Tensor, output: torch.Tensor = None):
|
||||
|
||||
"""
|
||||
Args:
|
||||
input: (..., 2*hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
Returns:
|
||||
output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
"""
|
||||
if output is None:
|
||||
output_shape = list(input.shape)
|
||||
output_shape[-1] = output_shape[-1] // 2
|
||||
output = input.new_empty(output_shape)
|
||||
|
||||
ops.infer.gelu_tanh_and_mul(input, output)
|
||||
|
||||
return output
|
||||
89
ixformer_sdk/inference/functions/act_bias_mm.py
Normal file
89
ixformer_sdk/inference/functions/act_bias_mm.py
Normal file
@@ -0,0 +1,89 @@
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
import torch.nn.functional as NNF
|
||||
|
||||
__all__ = ["act_bias_mm", "ref_act_bias_mm"]
|
||||
|
||||
|
||||
def ref_act_bias_mm(
|
||||
mat1: torch.Tensor,
|
||||
mat2: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
scale: float = 1,
|
||||
act_type: str = "none",
|
||||
trans_format: str = "NN",
|
||||
):
|
||||
assert len(mat1.shape) >= 2
|
||||
assert len(mat2.shape) >= 2
|
||||
if trans_format == "NN":
|
||||
if bias is not None:
|
||||
output = torch.matmul(mat1, mat2) * scale + bias
|
||||
else:
|
||||
output = torch.matmul(mat1, mat2) * scale
|
||||
else:
|
||||
if bias is not None:
|
||||
output = torch.matmul(mat1, mat2.transpose(-1, -2)) * scale + bias
|
||||
else:
|
||||
output = torch.matmul(mat1, mat2.transpose(-1, -2)) * scale
|
||||
if act_type == "gelu":
|
||||
output = NNF.gelu(output)
|
||||
elif act_type == "relu":
|
||||
output = NNF.relu(output)
|
||||
elif act_type == "silu":
|
||||
output = NNF.silu(output)
|
||||
elif act_type == "none":
|
||||
output = output
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
return output
|
||||
|
||||
|
||||
def act_bias_mm(
|
||||
mat1: torch.Tensor,
|
||||
mat2: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
output: torch.Tensor = None,
|
||||
scale: float = 1,
|
||||
act_type: str = "none",
|
||||
trans_format: str = "NN",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
mat1: [m,k] or [batch_count,m,k] torch.float16
|
||||
mat2: [k,n] or [n,k] torch.float16
|
||||
当trans_format为"NN"时[k,n], 当trans_format为"TN"时[n,k]
|
||||
bias: [n] torch.float16
|
||||
output: [m,n] torch.float16
|
||||
scale: float
|
||||
act_type: silu/gelu/relu/None str
|
||||
如果act_type不为None,则bias也不可以为None
|
||||
trans_format: NN or TN str
|
||||
Returns:
|
||||
output: [m,n] torch.float16
|
||||
"""
|
||||
assert len(mat1.shape) >= 2
|
||||
assert len(mat2.shape) >= 2
|
||||
if output is None:
|
||||
output_shape = list(mat1.shape)
|
||||
m = mat1.shape[-2]
|
||||
if trans_format == "NN":
|
||||
n = mat2.shape[-1]
|
||||
else:
|
||||
n = mat2.shape[-2]
|
||||
output_shape[-2] = m
|
||||
output_shape[-1] = n
|
||||
output = mat1.new_empty(output_shape)
|
||||
|
||||
add_bias = False
|
||||
if bias is not None:
|
||||
add_bias = True
|
||||
|
||||
if add_bias:
|
||||
ops.infer.act_bias_mm(
|
||||
mat1, mat2, bias, output, add_bias, scale, act_type, trans_format
|
||||
)
|
||||
else:
|
||||
ops.infer.act_bias_mm(
|
||||
mat1, mat2, mat1, output, add_bias, scale, act_type, trans_format
|
||||
)
|
||||
return output
|
||||
46
ixformer_sdk/inference/functions/add.py
Normal file
46
ixformer_sdk/inference/functions/add.py
Normal file
@@ -0,0 +1,46 @@
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"ref_add",
|
||||
"add",
|
||||
]
|
||||
|
||||
|
||||
def ref_add(input: torch.Tensor, other: torch.Tensor, out: torch.Tensor = None):
|
||||
return torch.add(input, other, out=out)
|
||||
|
||||
|
||||
def add(input: torch.Tensor, other: torch.Tensor, out: torch.Tensor = None):
|
||||
"""
|
||||
out = input + other
|
||||
Support elementwise addition, but broadcasting is not supported yet.
|
||||
Note: The dtype of input and other needs to be the same.
|
||||
Args:
|
||||
input: (...) torch.float32, torch.float16, torch.bfloat16
|
||||
other: (...) same as input
|
||||
out: (...) same as input
|
||||
Returns:
|
||||
out: (...) same as input
|
||||
"""
|
||||
if input.dtype not in [torch.float16, torch.float32, torch.bfloat16]:
|
||||
return torch.add(input, other, out=out)
|
||||
if not input.is_contiguous() or not other.is_contiguous():
|
||||
return torch.add(input, other, out=out)
|
||||
if out is not None and not out.is_contiguous():
|
||||
return torch.add(input, other, out=out)
|
||||
|
||||
if input.dtype != other.dtype:
|
||||
return torch.add(input, other, out=out)
|
||||
if out is not None and out.dtype != input.dtype:
|
||||
return torch.add(input, other, out=out)
|
||||
|
||||
assert input.shape == other.shape, (f"broadcasting is not supported yet."
|
||||
"input is {input.shape}, other is {other.shape}")
|
||||
|
||||
if out is None:
|
||||
out = torch.empty_like(input)
|
||||
|
||||
ops.infer.add(input, other, out)
|
||||
|
||||
return out
|
||||
199
ixformer_sdk/inference/functions/bert.py
Normal file
199
ixformer_sdk/inference/functions/bert.py
Normal file
@@ -0,0 +1,199 @@
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"ref_bert_embedding",
|
||||
"bert_embedding",
|
||||
"ref_bert_add_norm",
|
||||
"bert_add_norm",
|
||||
"ref_bert_unpack_start_end_logits",
|
||||
"bert_unpack_start_end_logits",
|
||||
"ref_bert_linear_residual",
|
||||
"bert_linear_residual",
|
||||
]
|
||||
|
||||
|
||||
def ref_bert_embedding(
|
||||
token_weight: torch.Tensor,
|
||||
pos_weight: torch.Tensor,
|
||||
type_weight: torch.Tensor,
|
||||
ln_weight: torch.Tensor,
|
||||
ln_bias: torch.Tensor,
|
||||
token_ids: torch.Tensor,
|
||||
pos_ids: torch.Tensor,
|
||||
type_ids: torch.Tensor,
|
||||
epsilon: float = 1e-5,
|
||||
out: torch.Tensor = None,
|
||||
):
|
||||
assert out is None
|
||||
emd1 = torch.nn.functional.embedding(token_ids, token_weight)
|
||||
emd2 = torch.nn.functional.embedding(pos_ids, pos_weight)
|
||||
emd3 = torch.nn.functional.embedding(type_ids, type_weight)
|
||||
|
||||
out = emd1 + emd2 + emd3
|
||||
out = torch.nn.functional.layer_norm(out, [out.shape[-1]], ln_weight, ln_bias)
|
||||
return out
|
||||
|
||||
|
||||
def bert_embedding(
|
||||
token_weight: torch.Tensor,
|
||||
pos_weight: torch.Tensor,
|
||||
type_weight: torch.Tensor,
|
||||
ln_weight: torch.Tensor,
|
||||
ln_bias: torch.Tensor,
|
||||
token_ids: torch.Tensor,
|
||||
pos_ids: torch.Tensor,
|
||||
type_ids: torch.Tensor,
|
||||
epsilon: float = 1e-5,
|
||||
out: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
token_weight: (vocab_size, hidden_size) torch.float16, torch.bfloat16
|
||||
pos_weight: (pos_size, hidden_size) same as token_weight
|
||||
type_weight: (type_size, hidden_size) same as token_weight
|
||||
ln_weight: (hidden_size) same as token_weight
|
||||
ln_bias: (hidden_size) same as token_weight
|
||||
token_ids: (num_tokens) torch.int32, torch.int64
|
||||
pos_ids: (num_tokens) same as token_ids
|
||||
type_ids: (num_tokens) same as token_ids
|
||||
epsilon: float
|
||||
out: (num_tokens, hidden_size) same as token_weight
|
||||
Returns:
|
||||
out: (num_tokens, hidden_size) same as token_weight
|
||||
"""
|
||||
if out is None:
|
||||
out_shape = list(token_ids.shape)
|
||||
hidden_size = token_weight.shape[-1]
|
||||
out_shape.append(hidden_size)
|
||||
out = token_weight.new_empty(out_shape)
|
||||
|
||||
ops.infer.bert_embedding(
|
||||
token_weight,
|
||||
pos_weight,
|
||||
type_weight,
|
||||
ln_weight,
|
||||
ln_bias,
|
||||
token_ids,
|
||||
pos_ids,
|
||||
type_ids,
|
||||
out,
|
||||
epsilon,
|
||||
)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def ref_bert_add_norm(
|
||||
input: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
ln_weight: torch.Tensor,
|
||||
ln_bias: torch.Tensor,
|
||||
epsilon: float = 1e-5,
|
||||
out: torch.Tensor = None,
|
||||
):
|
||||
assert out is None
|
||||
input = input + residual
|
||||
return torch.nn.functional.layer_norm(
|
||||
input, [input.shape[-1]], ln_weight, ln_bias, epsilon
|
||||
)
|
||||
|
||||
|
||||
def bert_add_norm(
|
||||
input: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
ln_weight: torch.Tensor,
|
||||
ln_bias: torch.Tensor,
|
||||
epsilon: float = 1e-5,
|
||||
out: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
out = input + residual
|
||||
out = add_norm(out, ln_weight, ln_bias, epsilon)
|
||||
Args:
|
||||
input: (num_tokens, hidden_size) torch.float16, torch.bfloat16
|
||||
residual: (num_tokens, hidden_size) same as input
|
||||
ln_weight: (hidden_size) same as input
|
||||
ln_bias: (hidden_size) same as input
|
||||
epsilon: float
|
||||
out: (num_tokens, hidden_size) same as input
|
||||
Returns:
|
||||
out: (num_tokens, hidden_size) same as input
|
||||
"""
|
||||
if out is None:
|
||||
out = torch.empty_like(input)
|
||||
ops.infer.bert_add_norm(input, residual, ln_weight, ln_bias, out, epsilon)
|
||||
return out
|
||||
|
||||
|
||||
def ref_bert_unpack_start_end_logits(
|
||||
logits: torch.Tensor,
|
||||
cu_seq_lens: torch.Tensor,
|
||||
max_seq_len: int,
|
||||
start_logits: torch.Tensor = None,
|
||||
end_logits: torch.Tensor = None,
|
||||
):
|
||||
batch_size = cu_seq_lens.shape[0] - 1
|
||||
if start_logits is None:
|
||||
start_logits = logits.new_empty([batch_size, max_seq_len])
|
||||
if end_logits is None:
|
||||
end_logits = logits.new_empty([batch_size, max_seq_len])
|
||||
cu_seq_len_cpu = cu_seq_lens.detach().cpu()
|
||||
for i in range(batch_size):
|
||||
start_idx = cu_seq_len_cpu[i]
|
||||
end_idx = cu_seq_len_cpu[i + 1]
|
||||
cur_len = end_idx - start_idx
|
||||
start_logits[i, :cur_len] = logits[start_idx:end_idx, 0]
|
||||
end_logits[i, :cur_len] = logits[start_idx:end_idx, 1]
|
||||
return start_logits, end_logits
|
||||
|
||||
|
||||
def bert_unpack_start_end_logits(
|
||||
logits: torch.Tensor,
|
||||
cu_seq_lens: torch.Tensor,
|
||||
max_seq_len: int,
|
||||
start_logits: torch.Tensor = None,
|
||||
end_logits: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
logits: (num_tokens, 2) torch.float16, torch.bfloat16
|
||||
cu_seq_lens: (batch_size+1) torch.int32, torch.int64
|
||||
max_seq_len: int
|
||||
start_logits: (batch_size, max_seq_len) same as logits
|
||||
end_logits: (batch_size, max_seq_len) same as logits
|
||||
Returns:
|
||||
start_logits: (batch_size, max_seq_len) same as logits
|
||||
end_logits: (batch_size, max_seq_len) same as logits
|
||||
"""
|
||||
batch_size = cu_seq_lens.shape[0] - 1
|
||||
if start_logits is None:
|
||||
start_logits = logits.new_empty([batch_size, max_seq_len])
|
||||
if end_logits is None:
|
||||
end_logits = logits.new_empty([batch_size, max_seq_len])
|
||||
ops.infer.bert_unpack_start_end_logits(
|
||||
logits, cu_seq_lens, start_logits, end_logits
|
||||
)
|
||||
return start_logits, end_logits
|
||||
|
||||
|
||||
def ref_bert_linear_residual(
|
||||
input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, out: torch.Tensor
|
||||
):
|
||||
return torch.nn.functional.linear(input, weight, bias) + out
|
||||
|
||||
|
||||
def bert_linear_residual(
|
||||
input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, out: torch.Tensor
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (m, k) torch.float16, torch.bfloat16
|
||||
weight: (n, k) same as input
|
||||
bias: (n) same as input
|
||||
out: (m, n) same as input
|
||||
Returns:
|
||||
out: (m, n) same as input
|
||||
"""
|
||||
ops.infer.bert_linear_residual(input, weight, bias, out)
|
||||
return out
|
||||
55
ixformer_sdk/inference/functions/bnb_dequant.py
Normal file
55
ixformer_sdk/inference/functions/bnb_dequant.py
Normal file
@@ -0,0 +1,55 @@
|
||||
from typing import List, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
from torch.autograd.function import Function, FunctionCtx
|
||||
|
||||
__all__ = [
|
||||
"bnb_dequant",
|
||||
"ref_bnb_dequant",
|
||||
]
|
||||
|
||||
|
||||
def ref_bnb_dequant(
|
||||
qA: torch.Tensor,
|
||||
SA: torch.Tensor,
|
||||
training: bool = False,
|
||||
scale: float = 127.0,
|
||||
dequant_type: int = 0,
|
||||
):
|
||||
A = torch.empty(qA.shape, dtype = SA.dtype, device = SA.device)
|
||||
if dequant_type == 0:
|
||||
for i in range(qA.size(0)):
|
||||
A[i:] = qA[i:] * (SA[i].to(torch.float) / scale).to(SA.dtype)
|
||||
else:
|
||||
for i in range(qA.size(1)):
|
||||
A[:,i] = qA[:,i] * (SA[i].to(torch.float) / scale).to(SA.dtype)
|
||||
|
||||
return A
|
||||
|
||||
|
||||
def bnb_dequant(
|
||||
qA: torch.Tensor,
|
||||
SA: torch.Tensor,
|
||||
training: bool = False,
|
||||
scale: float = 127.0,
|
||||
dequant_type: int = 0,
|
||||
) -> torch.Tensor:
|
||||
|
||||
"""
|
||||
Args:
|
||||
qA: (row, col) torch.int8
|
||||
dequant input
|
||||
SA: (row) or (col) torch.half
|
||||
scale vector
|
||||
training: bool
|
||||
scale: float
|
||||
dequnt_type: int
|
||||
0 : every row shared a scale, SA shape : [row]
|
||||
1 : every col shared a scale, SA shape : [col]
|
||||
Returns:
|
||||
Tensor: (row, col) torch.half
|
||||
dequant output
|
||||
|
||||
"""
|
||||
return ops.infer.bnb_dequant(qA, SA, scale, dequant_type)
|
||||
175
ixformer_sdk/inference/functions/bnb_double_quant.py
Normal file
175
ixformer_sdk/inference/functions/bnb_double_quant.py
Normal file
@@ -0,0 +1,175 @@
|
||||
from typing import List, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
from torch.autograd.function import Function, FunctionCtx
|
||||
|
||||
__all__ = ["bnb_double_quant"]
|
||||
|
||||
import ctypes as ct
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
def get_ptr(A):
|
||||
if A is None:
|
||||
return None
|
||||
else:
|
||||
return ct.c_void_p(A.data.data_ptr())
|
||||
|
||||
|
||||
class COOSparseTensor:
|
||||
def __init__(self, rows, cols, nnz, rowidx, colidx, values):
|
||||
assert rowidx.dtype == torch.int
|
||||
assert colidx.dtype == torch.int
|
||||
assert values.dtype == torch.half
|
||||
assert values.numel() == nnz
|
||||
assert rowidx.numel() == nnz
|
||||
assert colidx.numel() == nnz
|
||||
|
||||
self.rows = rows
|
||||
self.cols = cols
|
||||
self.nnz = nnz
|
||||
self.rowidx = rowidx
|
||||
self.colidx = colidx
|
||||
self.values = values
|
||||
|
||||
|
||||
def coo_zeros(rows, cols, nnz, device, dtype=torch.half):
|
||||
rowidx = torch.full(size=(nnz,), fill_value=0, dtype=torch.int, device=device)
|
||||
|
||||
colidx = torch.full((nnz,), fill_value=0, dtype=torch.int, device=device)
|
||||
values = torch.full((nnz,), fill_value=0, dtype=dtype, device=device)
|
||||
return COOSparseTensor(rows, cols, nnz, rowidx, colidx, values)
|
||||
|
||||
|
||||
def get_colrow_absmax(
|
||||
A, row_stats=None, col_stats=None, nnz_block_ptr=None, threshold=0.0
|
||||
):
|
||||
cols = A.shape[-1]
|
||||
if len(A.shape) == 3:
|
||||
rows = A.shape[0] * A.shape[1]
|
||||
else:
|
||||
rows = A.shape[0]
|
||||
|
||||
col_tiles = (cols + 255) // 256
|
||||
tiled_rows = ((rows + 15) // 16) * 16
|
||||
if row_stats is None:
|
||||
row_stats = torch.full(
|
||||
size=(rows,), fill_value=-50000.0, dtype=torch.float, device=A.device
|
||||
)
|
||||
if col_stats is None:
|
||||
col_stats = torch.full(
|
||||
size=(cols,), fill_value=-50000.0, dtype=torch.float, device=A.device
|
||||
)
|
||||
|
||||
# if nnz_block_ptr is None and threshold > 0.0:
|
||||
nnz_block_ptr = torch.full(
|
||||
size=(tiled_rows * col_tiles + 1,),
|
||||
fill_value=0,
|
||||
dtype=torch.int,
|
||||
device=A.device,
|
||||
)
|
||||
|
||||
ops.infer.bnb_getColRowStats(
|
||||
A, row_stats, col_stats, nnz_block_ptr, threshold, rows, cols
|
||||
)
|
||||
|
||||
return row_stats, col_stats, nnz_block_ptr
|
||||
|
||||
|
||||
# A : quant input shape : [row, col] shape:torch.half
|
||||
def bnb_double_quant(
|
||||
A: torch.Tensor, training: bool = False, threshold: float = 0.0
|
||||
) -> torch.Tensor:
|
||||
|
||||
"""
|
||||
Args:
|
||||
A: (row, col) torch.float16
|
||||
quant input
|
||||
training: bool
|
||||
threshold: float
|
||||
abs of element exceeds threshold will be ignored
|
||||
Returns:
|
||||
out_row: (row, col) torch.int8
|
||||
out_col: (row, col) torch.int8
|
||||
row_stats (row) torch.float
|
||||
col_stats (col) torch.float
|
||||
coo_tensor
|
||||
"""
|
||||
|
||||
assert A.dtype == torch.half
|
||||
|
||||
cols = A.shape[-1]
|
||||
if len(A.shape) == 3:
|
||||
rows = A.shape[0] * A.shape[1]
|
||||
else:
|
||||
rows = A.shape[0]
|
||||
|
||||
row_stats, col_stats, nnz_row_ptr = get_colrow_absmax(A, threshold=threshold)
|
||||
|
||||
out_col = torch.full(size=A.shape, fill_value=0, dtype=torch.int8, device=A.device)
|
||||
out_row = torch.full(size=A.shape, fill_value=0, dtype=torch.int8, device=A.device)
|
||||
|
||||
coo_tensor = None
|
||||
if threshold > 0.0:
|
||||
nnz = nnz_row_ptr.cpu().numpy()[-1]
|
||||
if nnz > 0:
|
||||
coo_tensor = coo_zeros(A.shape[0], A.shape[1], nnz, A.device)
|
||||
|
||||
ops.infer.bnb_doubleRowColQuant(
|
||||
A,
|
||||
row_stats,
|
||||
col_stats,
|
||||
out_col,
|
||||
out_row,
|
||||
coo_tensor.rowidx,
|
||||
coo_tensor.colidx,
|
||||
coo_tensor.values,
|
||||
nnz_row_ptr,
|
||||
threshold,
|
||||
rows,
|
||||
cols,
|
||||
)
|
||||
val, idx = torch.sort(torch.Tensor(coo_tensor.rowidx.cpu().numpy()))
|
||||
coo_tensor.rowidx = val
|
||||
coo_tensor.colidx = torch.Tensor(coo_tensor.colidx.cpu().numpy())[idx].to(
|
||||
torch.int32
|
||||
)
|
||||
coo_tensor.values = torch.Tensor(coo_tensor.values.cpu().numpy())[idx].to(
|
||||
torch.half
|
||||
)
|
||||
# coo_tensor.colidx = coo_tensor.colidx[idx]
|
||||
# coo_tensor.values = coo_tensor.values[idx]
|
||||
else:
|
||||
ops.infer.bnb_doubleRowColQuant(
|
||||
A,
|
||||
row_stats,
|
||||
col_stats,
|
||||
out_col,
|
||||
out_row,
|
||||
out_row,
|
||||
out_row,
|
||||
out_row,
|
||||
out_row,
|
||||
0.0,
|
||||
rows,
|
||||
cols,
|
||||
)
|
||||
else:
|
||||
ops.infer.bnb_doubleRowColQuant(
|
||||
A,
|
||||
row_stats,
|
||||
col_stats,
|
||||
out_col,
|
||||
out_row,
|
||||
out_row,
|
||||
out_row,
|
||||
out_row,
|
||||
out_row,
|
||||
threshold,
|
||||
rows,
|
||||
cols,
|
||||
)
|
||||
|
||||
return out_row, out_col, row_stats, col_stats, coo_tensor
|
||||
73
ixformer_sdk/inference/functions/bnb_mm_dequant.py
Normal file
73
ixformer_sdk/inference/functions/bnb_mm_dequant.py
Normal file
@@ -0,0 +1,73 @@
|
||||
from typing import List, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
from torch.autograd.function import Function, FunctionCtx
|
||||
|
||||
__all__ = ["bnb_mm_dequant"]
|
||||
|
||||
|
||||
# A : quant input shape : [row, col] shape : torch.int
|
||||
def bnb_mm_dequant(
|
||||
A: torch.Tensor,
|
||||
quant_state: tuple,
|
||||
row_stats: torch.Tensor,
|
||||
col_stats: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
add_bias: bool = False,
|
||||
training: bool = False,
|
||||
) -> torch.Tensor:
|
||||
|
||||
"""
|
||||
Args:
|
||||
A: (row, col) torch.int8
|
||||
quant_state: tuple
|
||||
row_stats: (row) torch.float
|
||||
col_stats: (col) torch.float
|
||||
bias: (col) torch.half
|
||||
add_bias: bool
|
||||
training: bool
|
||||
Returns:
|
||||
Tensor: (row, col) torch.half
|
||||
|
||||
"""
|
||||
|
||||
assert A.dtype == torch.int
|
||||
if bias is not None:
|
||||
add_bias = True
|
||||
print("bias.dtype:", bias.dtype)
|
||||
assert bias.dtype == torch.half
|
||||
else:
|
||||
bias = A
|
||||
out_shape = quant_state[0]
|
||||
if len(out_shape) == 3:
|
||||
out_shape = (out_shape[0] * out_shape[1], out_shape[2])
|
||||
out = torch.full(size=out_shape, fill_value=0, dtype=torch.half, device=A.device)
|
||||
new_row_stats = torch.full(
|
||||
size=(out_shape[0],), fill_value=0, dtype=torch.float, device=A.device
|
||||
)
|
||||
new_col_stats = torch.full(
|
||||
size=(out_shape[1],), fill_value=0, dtype=torch.float, device=A.device
|
||||
)
|
||||
|
||||
assert (
|
||||
new_row_stats.shape[0] == row_stats.shape[0]
|
||||
), f"{new_row_stats.shape} vs {row_stats.shape}"
|
||||
assert (
|
||||
new_col_stats.shape[0] == col_stats.shape[0]
|
||||
), f"{new_col_stats.shape} vs {col_stats.shape}"
|
||||
numRows = out_shape[0]
|
||||
numCols = out_shape[1]
|
||||
ops.infer.bnb_mm_dequant(
|
||||
A,
|
||||
row_stats,
|
||||
col_stats,
|
||||
out,
|
||||
new_row_stats,
|
||||
new_col_stats,
|
||||
numRows,
|
||||
numCols,
|
||||
add_bias,
|
||||
bias,
|
||||
)
|
||||
return out
|
||||
56
ixformer_sdk/inference/functions/bnb_qgemm.py
Normal file
56
ixformer_sdk/inference/functions/bnb_qgemm.py
Normal file
@@ -0,0 +1,56 @@
|
||||
from typing import List, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = ["bnb_qgemm", "ref_bnb_qgemm"]
|
||||
|
||||
|
||||
# qA : quant input shape : [bs, in_feature]
|
||||
# qW : quant weight shape : [out_feature, in_feature]
|
||||
# SA : scale vector of qA shape : [bs]
|
||||
# SW : scale vector of qW shape : [out_feature]
|
||||
|
||||
|
||||
def ref_bnb_qgemm(
|
||||
qA: torch.Tensor,
|
||||
qW: torch.Tensor,
|
||||
SA: torch.Tensor,
|
||||
SW: torch.Tensor,
|
||||
training: bool = False,
|
||||
scaleA: float = 127.0,
|
||||
scaleW: float = 127.0,
|
||||
):
|
||||
y = torch.nn.functional.linear(qA.to(torch.float), qW.to(torch.float))
|
||||
out = torch.empty(y.shape, dtype = SA.dtype, device = SA.device)
|
||||
for i in range(qA.size(0)):
|
||||
for j in range(qW.size(0)):
|
||||
out[i][j] = y[i][j] * (SA[i].to(torch.float) / scaleA) * (SW[j].to(torch.float) / scaleW)
|
||||
return out.to(SA.dtype)
|
||||
|
||||
|
||||
def bnb_qgemm(
|
||||
qA: torch.Tensor,
|
||||
qW: torch.Tensor,
|
||||
SA: torch.Tensor,
|
||||
SW: torch.Tensor,
|
||||
training: bool = False,
|
||||
scaleA: float = 127.0,
|
||||
scaleW: float = 127.0,
|
||||
) -> torch.Tensor:
|
||||
|
||||
"""
|
||||
Args:
|
||||
qA: (bs, in_feature) torch.int8
|
||||
qW: (out_feature, in_feature) torch.int8
|
||||
SA: (bs) torch.half
|
||||
scale vector of qA
|
||||
SA: (out_feature) torch.half
|
||||
scale vector of qW
|
||||
training: bool
|
||||
scaleA: float
|
||||
scaleW: float
|
||||
Returns:
|
||||
Tensor: (bs, out_feature) torch.half
|
||||
"""
|
||||
return ops.infer.bnb_qgemm(qA, qW, SA, SW, scaleA, scaleW)
|
||||
58
ixformer_sdk/inference/functions/bnb_quant.py
Normal file
58
ixformer_sdk/inference/functions/bnb_quant.py
Normal file
@@ -0,0 +1,58 @@
|
||||
from typing import List, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = ["bnb_quant", "ref_bnb_quant"]
|
||||
|
||||
|
||||
# A : input shape : [row, col]
|
||||
# SA : scale vector
|
||||
# quant_type
|
||||
# 0 : every row shared a scale, SA shape : [row]
|
||||
# 1 : every col shared a scale, SA shape : [col]
|
||||
def ref_bnb_quant(
|
||||
A: torch.Tensor,
|
||||
SA: torch.Tensor,
|
||||
training: bool = False,
|
||||
scale: float = 127.0,
|
||||
quant_type: int = 0,
|
||||
):
|
||||
qA = torch.empty(A.shape, device = SA.device)
|
||||
if quant_type == 0:
|
||||
for i in range(A.size(0)):
|
||||
qA[i:] = torch.round(A[i:] * (scale / SA[i].to(torch.float)))
|
||||
else:
|
||||
for i in range(A.size(1)):
|
||||
qA[:,i] = torch.round(A[:,i] * (scale / SA[i].to(torch.float)))
|
||||
|
||||
qA_clamped = torch.clamp(qA, min=-128, max=127)
|
||||
qA = qA_clamped.to(torch.int8)
|
||||
return qA
|
||||
|
||||
|
||||
def bnb_quant(
|
||||
A: torch.Tensor,
|
||||
SA: torch.Tensor,
|
||||
training: bool = False,
|
||||
scale: float = 127.0,
|
||||
quant_type: int = 0,
|
||||
) -> torch.Tensor:
|
||||
|
||||
"""
|
||||
Args:
|
||||
A: (row, col) torch.half
|
||||
quant input
|
||||
SA: (row) or (col) torch.half
|
||||
scale vector
|
||||
training: bool
|
||||
scale: float
|
||||
qunt_type: int
|
||||
0 : every row shared a scale, SA shape : [row]
|
||||
1 : every col shared a scale, SA shape : [col]
|
||||
Returns:
|
||||
Tensor: (row, col) torch.int8
|
||||
quant output
|
||||
|
||||
"""
|
||||
return ops.infer.bnb_quant(A, SA, scale, quant_type)
|
||||
52
ixformer_sdk/inference/functions/bnb_rowcol_absmax.py
Normal file
52
ixformer_sdk/inference/functions/bnb_rowcol_absmax.py
Normal file
@@ -0,0 +1,52 @@
|
||||
from typing import List, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = ["bnb_rowcol_absmax", "ref_bnb_rowcol_absmax"]
|
||||
|
||||
|
||||
# input : input shape : [row, col]
|
||||
# threshold : abs of element exceeds threshold will be ignored
|
||||
# type
|
||||
# 0 : row absmax
|
||||
def ref_bnb_rowcol_absmax(
|
||||
input: torch.Tensor,
|
||||
training: bool = False,
|
||||
threshold: float = 0.0,
|
||||
type: int = 0,
|
||||
):
|
||||
input = input.float()
|
||||
if threshold ==0.0:
|
||||
threshold = float('inf')
|
||||
mask = (torch.abs(input) < threshold)
|
||||
masked_input = mask * input
|
||||
masked_input = masked_input.half()
|
||||
if type == 0:
|
||||
out = torch.amax(torch.abs(masked_input), dim=1)
|
||||
|
||||
else:
|
||||
out = torch.amax(torch.abs(masked_input), dim=0)
|
||||
return out
|
||||
|
||||
|
||||
def bnb_rowcol_absmax(
|
||||
input: torch.Tensor,
|
||||
training: bool = False,
|
||||
threshold: float = 0.0,
|
||||
type: int = 0,
|
||||
) -> torch.Tensor:
|
||||
|
||||
"""
|
||||
Args:
|
||||
input: (row, col) torch.half
|
||||
目前col值必须满足col%2==0
|
||||
training: bool
|
||||
threshold: float
|
||||
abs of element exceeds threshold will be ignored
|
||||
type: int
|
||||
row absmax, 目前只支持type=0
|
||||
Returns:
|
||||
Tensor: (row) torch.half
|
||||
"""
|
||||
return ops.infer.bnb_rowcol_absmax(input, threshold, type)
|
||||
198
ixformer_sdk/inference/functions/conv2d.py
Normal file
198
ixformer_sdk/inference/functions/conv2d.py
Normal file
@@ -0,0 +1,198 @@
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
import torch.nn.functional as NNF
|
||||
|
||||
__all__ = ["conv2d", "ref_conv2d", "ref_conv2d_nhwc", "conv2d_nhwc"]
|
||||
|
||||
|
||||
def is_channels_last(ten):
|
||||
return torch._prims_common.suggest_memory_format(ten) == torch.channels_last
|
||||
|
||||
|
||||
def _pair(x):
|
||||
if isinstance(x, (list, tuple)):
|
||||
return x
|
||||
return (x, x)
|
||||
|
||||
|
||||
def ref_conv2d(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
stride: Union[int, tuple] = 1,
|
||||
padding: Union[int, tuple] = 0,
|
||||
dilation: Union[int, tuple] = 1,
|
||||
groups: int = 1,
|
||||
):
|
||||
|
||||
output = NNF.conv2d(input, weight, bias, stride, padding, dilation, groups)
|
||||
return output
|
||||
|
||||
|
||||
# conv2d官方接口，如果weight是torch.channels_last,输出也是torch.channels_last；如果weight是nchw，那么输出也是nchw;特殊情况，如果输入是nchw，weight是torch.channels_last，输出也是torch.channels_last
|
||||
def conv2d(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
stride: Union[int, tuple] = 1,
|
||||
padding: Union[int, tuple] = 0,
|
||||
dilation: Union[int, tuple] = 1,
|
||||
groups: int = 1,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (n,in_c,h,w) torch.float16
|
||||
weight: (out_c,in_c/groups,kH,kW) torch.float16
|
||||
bias: (out_c) torch.float16
|
||||
stride: int or tuple
|
||||
Stride of the convolution. Default: 1
|
||||
padding: int or tuple
|
||||
Padding added to all four sides of the input. Default: 0
|
||||
dilation: int or tuple
|
||||
Spacing between kernel elements. Default: 1
|
||||
groups: int
|
||||
Number of blocked connections from input channels to output channels. Default: 1
|
||||
Returns:
|
||||
Tensor: (n,out_c,h_out,w_out) torch.float16
|
||||
h_out = (h_in + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) / stride_h + 1;
|
||||
w_out = (w_in + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) / stride_w + 1;
|
||||
"""
|
||||
stride = _pair(stride)
|
||||
padding = _pair(padding)
|
||||
dilation = _pair(dilation)
|
||||
|
||||
channel_last = is_channels_last(weight)
|
||||
if not is_channels_last(input) and channel_last:
|
||||
input = input.to(memory_format=torch.channels_last)
|
||||
|
||||
# compute outshape
|
||||
n, in_c, h_in, w_in = input.shape
|
||||
out_c, _, kernel_h, kernel_w = weight.shape
|
||||
pad_h = padding[0]
|
||||
pad_w = padding[1]
|
||||
stride_h = stride[0]
|
||||
stride_w = stride[1]
|
||||
dilation_h = dilation[0]
|
||||
dilation_w = dilation[1]
|
||||
h_out = (h_in + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) // stride_h + 1
|
||||
w_out = (w_in + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) // stride_w + 1
|
||||
|
||||
if channel_last:
|
||||
output_shape = [n, out_c, h_out, w_out]
|
||||
output = torch.empty(
|
||||
output_shape,
|
||||
memory_format=torch.channels_last,
|
||||
dtype=input.dtype,
|
||||
device=input.device,
|
||||
)
|
||||
else:
|
||||
output_shape = [n, out_c, h_out, w_out]
|
||||
output = input.new_empty(output_shape)
|
||||
|
||||
if channel_last:
|
||||
input = input.permute(0, 2, 3, 1)
|
||||
weight = weight.permute(0, 2, 3, 1)
|
||||
output = output.permute(0, 2, 3, 1)
|
||||
|
||||
if bias is not None:
|
||||
bias = bias.float()
|
||||
ops.infer.conv2d(
|
||||
input, weight, bias, output, stride, padding, dilation, groups, channel_last
|
||||
)
|
||||
if channel_last:
|
||||
output = output.permute(0, 3, 1, 2)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def ref_conv2d_nhwc(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
stride: Union[int, tuple] = 1,
|
||||
padding: Union[int, tuple] = 0,
|
||||
dilation: Union[int, tuple] = 1,
|
||||
groups: int = 1,
|
||||
):
|
||||
|
||||
output = NNF.conv2d(
|
||||
input.permute(0, 3, 1, 2).contiguous(),
|
||||
weight.permute(0, 3, 1, 2).contiguous(),
|
||||
bias,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
)
|
||||
return output.permute(0, 2, 3, 1).contiguous()
|
||||
|
||||
|
||||
# conv2d_nhwc，
|
||||
# conv2d官方接口解决两种情况：
|
||||
# 1、务必输入tensor内存上是nhwc,且tensor属于memory_format=torch.channels_last，
|
||||
# 2、或者输入tensor内存上是nchw，并且是contiguous；
|
||||
# conv2d官方接口不能解决，conv2d_nhwc则可处理这种情况的
|
||||
# 输入tensor内存上是nhwc的，但tensor没有用memory_format=torch.channels_last进行过处理，不会有memory_format=torch.channels_last的标签
|
||||
def conv2d_nhwc(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
stride: Union[int, tuple] = 1,
|
||||
padding: Union[int, tuple] = 0,
|
||||
dilation: Union[int, tuple] = 1,
|
||||
groups: int = 1,
|
||||
):
|
||||
|
||||
"""
|
||||
Args:
|
||||
input: (n,h,w,in_c) torch.float16
|
||||
weight: (out_c,kH,kW,in_c/groups) torch.float16
|
||||
bias: (out_c) torch.float16
|
||||
stride: int or tuple
|
||||
Stride of the convolution. Default: 1
|
||||
padding: int or tuple
|
||||
Padding added to all four sides of the input. Default: 0
|
||||
dilation: int or tuple
|
||||
Spacing between kernel elements. Default: 1
|
||||
groups: int
|
||||
Number of blocked connections from input channels to output channels. Default: 1
|
||||
Returns:
|
||||
Tensor: (n,h_out,w_out,out_c) torch.float16
|
||||
h_out = (h_in + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) / stride_h + 1;
|
||||
w_out = (w_in + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) / stride_w + 1;
|
||||
"""
|
||||
|
||||
stride = _pair(stride)
|
||||
padding = _pair(padding)
|
||||
dilation = _pair(dilation)
|
||||
|
||||
assert input.is_contiguous()
|
||||
assert weight.is_contiguous()
|
||||
|
||||
# compute outshape
|
||||
n, h_in, w_in, in_c = input.shape
|
||||
(
|
||||
out_c,
|
||||
kernel_h,
|
||||
kernel_w,
|
||||
_,
|
||||
) = weight.shape
|
||||
pad_h = padding[0]
|
||||
pad_w = padding[1]
|
||||
stride_h = stride[0]
|
||||
stride_w = stride[1]
|
||||
dilation_h = dilation[0]
|
||||
dilation_w = dilation[1]
|
||||
h_out = (h_in + 2 * pad_h - dilation_h * (kernel_h - 1) - 1) // stride_h + 1
|
||||
w_out = (w_in + 2 * pad_w - dilation_w * (kernel_w - 1) - 1) // stride_w + 1
|
||||
|
||||
output_shape = [n, h_out, w_out, out_c]
|
||||
output = torch.empty(output_shape, dtype=input.dtype, device=input.device)
|
||||
if bias is not None:
|
||||
bias = bias.float()
|
||||
ops.infer.conv2d(
|
||||
input, weight, bias, output, stride, padding, dilation, groups, True
|
||||
)
|
||||
return output
|
||||
204
ixformer_sdk/inference/functions/cross_entropy_loss.py
Normal file
204
ixformer_sdk/inference/functions/cross_entropy_loss.py
Normal file
@@ -0,0 +1,204 @@
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = ["vocab_parallel_cross_entropy", "ref_vocab_parallel_cross_entropy"]
|
||||
|
||||
|
||||
def ref_vocab_parallel_cross_entropy(
|
||||
vocab_parallel_logits: torch.Tensor,
|
||||
target: torch.Tensor,
|
||||
label_smoothing: float = 0.0,
|
||||
world_size: int = 1,
|
||||
vocab_start_index: int = 0,
|
||||
vocab_end_index: int = 320000,
|
||||
group=None,
|
||||
):
|
||||
if world_size == 1:
|
||||
vocab_parallel_logits = vocab_parallel_logits.float()
|
||||
partition_vocab_size = vocab_parallel_logits.size()[-1]
|
||||
logits_max = torch.max(vocab_parallel_logits, dim=-1)[0]
|
||||
vocab_parallel_logits = vocab_parallel_logits - logits_max.unsqueeze(dim=-1)
|
||||
target_mask = (target < vocab_start_index) | (target >= vocab_end_index)
|
||||
masked_target = target.clone() - vocab_start_index
|
||||
masked_target[target_mask] = 0
|
||||
logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size)
|
||||
masked_target_1d = masked_target.view(-1)
|
||||
arange_1d = torch.arange(
|
||||
start=0, end=logits_2d.size()[0], device=logits_2d.device
|
||||
)
|
||||
|
||||
predicted_logits_1d = logits_2d[arange_1d, masked_target_1d]
|
||||
predicted_logits_1d = predicted_logits_1d.clone().contiguous()
|
||||
|
||||
predicted_logits = predicted_logits_1d.view_as(target)
|
||||
predicted_logits[target_mask] = 0.0
|
||||
exp_logits = vocab_parallel_logits
|
||||
torch.exp(vocab_parallel_logits, out=exp_logits)
|
||||
sum_exp_logits = exp_logits.sum(dim=-1)
|
||||
loss = torch.log(sum_exp_logits) - predicted_logits
|
||||
exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1))
|
||||
if label_smoothing > 0:
|
||||
"""
|
||||
We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth.
|
||||
= (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt})
|
||||
= (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
|
||||
= ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
|
||||
= (K * (1 - alpha) - 1) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i} y_i
|
||||
= (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K
|
||||
From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py
|
||||
"""
|
||||
assert 1.0 > label_smoothing > 0.0
|
||||
smoothing = label_smoothing * partition_vocab_size / (partition_vocab_size - 1)
|
||||
|
||||
# Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs.
|
||||
log_probs = torch.log(exp_logits)
|
||||
mean_log_probs = log_probs.mean(dim=-1)
|
||||
loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs
|
||||
|
||||
else:
|
||||
# Maximum value along vocab dimension across all GPUs.
|
||||
logits_max = torch.max(vocab_parallel_logits, dim=-1)[0]
|
||||
torch.distributed.all_reduce(
|
||||
logits_max, op=torch.distributed.ReduceOp.MAX, group=group
|
||||
)
|
||||
# Subtract the maximum value.
|
||||
vocab_parallel_logits = vocab_parallel_logits - logits_max.unsqueeze(dim=-1)
|
||||
|
||||
# Get the partition's vocab indecies
|
||||
partition_vocab_size = vocab_parallel_logits.size()[-1]
|
||||
|
||||
# Create a mask of valid vocab ids (1 means it needs to be masked).
|
||||
target_mask = (target < vocab_start_index) | (target >= vocab_end_index)
|
||||
masked_target = target.clone() - vocab_start_index
|
||||
masked_target[target_mask] = 0
|
||||
|
||||
# Get predicted-logits = logits[target].
|
||||
# For Simplicity, we convert logits to a 2-D tensor with size
|
||||
# [*, partition-vocab-size] and target to a 1-D tensor of size [*].
|
||||
logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size)
|
||||
masked_target_1d = masked_target.view(-1)
|
||||
arange_1d = torch.arange(
|
||||
start=0, end=logits_2d.size()[0], device=logits_2d.device
|
||||
)
|
||||
predicted_logits_1d = logits_2d[arange_1d, masked_target_1d]
|
||||
predicted_logits_1d = predicted_logits_1d.clone().contiguous()
|
||||
predicted_logits = predicted_logits_1d.view_as(target)
|
||||
predicted_logits[target_mask] = 0.0
|
||||
# All reduce is needed to get the chunks from other GPUs.
|
||||
torch.distributed.all_reduce(
|
||||
predicted_logits,
|
||||
op=torch.distributed.ReduceOp.SUM,
|
||||
group=group,
|
||||
)
|
||||
|
||||
# Sum of exponential of logits along vocab dimension across all GPUs.
|
||||
exp_logits = vocab_parallel_logits
|
||||
torch.exp(vocab_parallel_logits, out=exp_logits)
|
||||
sum_exp_logits = exp_logits.sum(dim=-1)
|
||||
torch.distributed.all_reduce(
|
||||
sum_exp_logits,
|
||||
op=torch.distributed.ReduceOp.SUM,
|
||||
group=group,
|
||||
)
|
||||
|
||||
# Loss = log(sum(exp(logits))) - predicted-logit.
|
||||
loss = torch.log(sum_exp_logits) - predicted_logits
|
||||
|
||||
# Normalize and optionally smooth logits
|
||||
exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1))
|
||||
|
||||
vocab_size = exp_logits.size(-1)
|
||||
if label_smoothing > 0:
|
||||
"""
|
||||
We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth.
|
||||
= (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt})
|
||||
= (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
|
||||
= ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
|
||||
= (K * (1 - alpha) - 1) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i} y_i
|
||||
= (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K
|
||||
From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py
|
||||
"""
|
||||
assert 1.0 > label_smoothing > 0.0
|
||||
smoothing = label_smoothing * vocab_size / (vocab_size - 1)
|
||||
|
||||
# Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs.
|
||||
log_probs = torch.log(exp_logits)
|
||||
mean_log_probs = log_probs.mean(dim=-1)
|
||||
loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
def vocab_parallel_cross_entropy(
|
||||
vocab_parallel_logits: torch.Tensor,
|
||||
target: torch.Tensor,
|
||||
label_smoothing: float = 0.0,
|
||||
world_size: int = 1,
|
||||
vocab_start_index: int = 0,
|
||||
vocab_end_index: int = 320000,
|
||||
group=None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
vocab_parallel_logits: (seq_len,1,vocal_size) torch.float16, torch.bfloat16, torch.float
|
||||
target: (seq_len,1) torch.int64
|
||||
label_smoothing: float
|
||||
默认为0.0. 用于标签平滑
|
||||
world_size: int
|
||||
当world_size = 1时,目前只支持batch_size = 1 的情况
|
||||
vocab_start_index: int
|
||||
vocab_end_index: int
|
||||
group:
|
||||
TP 并行组
|
||||
Returns:
|
||||
loss: (seq_len,1) torch.float
|
||||
"""
|
||||
if world_size == 1:
|
||||
device = vocab_parallel_logits.device
|
||||
xnumel = vocab_parallel_logits.shape[0]
|
||||
rnumel = vocab_parallel_logits.shape[-1]
|
||||
exp_logits = torch.empty(
|
||||
(xnumel, 1, rnumel), device=device, dtype=torch.float32
|
||||
)
|
||||
masked_target_1d = torch.empty((xnumel,), device=device, dtype=torch.int32)
|
||||
loss = torch.empty((xnumel, 1), device=device, dtype=torch.float32)
|
||||
|
||||
ops.train.cross_entropy_loss_forward(
|
||||
vocab_parallel_logits, target.int(), exp_logits, masked_target_1d, loss
|
||||
)
|
||||
|
||||
vocab_size = exp_logits.size(-1)
|
||||
if label_smoothing > 0:
|
||||
"""
|
||||
We'd like to assign 1 / (K - 1) probability mass to every index that is not the ground truth.
|
||||
= (1 - alpha) * y_gt + alpha * mean(y_{i for i != gt})
|
||||
= (1 - alpha) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
|
||||
= ((K - 1) * (1 - alpha) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i != gt} y_i
|
||||
= (K * (1 - alpha) - 1) / (K - 1)) * y_gt + (alpha / (K - 1)) * \sum_{i} y_i
|
||||
= (1 - (alpha * K) / (K - 1)) * y_gt + ( (alpha * K) / (K - 1) ) * \sum_{i} y_i / K
|
||||
From: https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/losses/smoothed_cross_entropy.py
|
||||
"""
|
||||
assert 1.0 > label_smoothing > 0.0
|
||||
smoothing = label_smoothing * vocab_size / (vocab_size - 1)
|
||||
|
||||
# Exp logits at this point are normalized probabilities. So we can just take the log to get log-probs.
|
||||
log_probs = torch.log(exp_logits)
|
||||
mean_log_probs = log_probs.mean(dim=-1)
|
||||
loss = (1.0 - smoothing) * loss - smoothing * mean_log_probs
|
||||
|
||||
# Store softmax, target-mask and masked-target for backward pass.
|
||||
|
||||
return loss
|
||||
else:
|
||||
loss = ref_vocab_parallel_cross_entropy(
|
||||
vocab_parallel_logits,
|
||||
target,
|
||||
label_smoothing,
|
||||
world_size,
|
||||
vocab_start_index,
|
||||
vocab_end_index,
|
||||
group,
|
||||
)
|
||||
return loss
|
||||
201
ixformer_sdk/inference/functions/flash_attn.py
Normal file
201
ixformer_sdk/inference/functions/flash_attn.py
Normal file
@@ -0,0 +1,201 @@
|
||||
import math
|
||||
from typing import List, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
from torch.autograd.function import Function, FunctionCtx
|
||||
|
||||
__all__ = [
|
||||
"ixinfer_flash_attn_unpad",
|
||||
"ixinfer_flash_attn_pad",
|
||||
"ref_ixinfer_flash_attn_pad",
|
||||
]
|
||||
|
||||
|
||||
def ixinfer_flash_attn_unpad(
|
||||
# total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
|
||||
q: "torch.Tensor",
|
||||
# total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
|
||||
k: "torch.Tensor",
|
||||
# total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
|
||||
v: "torch.Tensor",
|
||||
# total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
|
||||
cu_seqlens_q: "torch.Tensor", # b+1
|
||||
cu_seqlens_k: "torch.Tensor", # b+1
|
||||
max_seqlen_q: int,
|
||||
max_seqlen_k: int,
|
||||
is_causal: bool = False,
|
||||
atten_scale: float = None,
|
||||
sqrt_alibi: bool = False,
|
||||
# total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
|
||||
alibi_slopes: "torch.Tensor" = None,
|
||||
out: "torch.Tenosr" = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
q: (total_q, nheads, headdim) torch.float16, torch.bfloat16
|
||||
where total_q = total number of query tokens in the batch.
|
||||
k: (total_k, nheads_k, headdim) torch.float16, torch.bfloat16
|
||||
where total_k = total number of key tokens in the batch.
|
||||
v: (total_k, nheads_k, headdim) torch.float16, torch.bfloat16
|
||||
cu_seqlens_q: (batch_size + 1) torch.int32
|
||||
The cumulative sequence lengths of the sequences in the batch, used to index into q.
|
||||
cu_seqlens_k: (batch_size + 1) torch.int32
|
||||
The cumulative sequence lengths of the sequences in the batch, used to index into kv.
|
||||
max_seqlen_q: int
|
||||
Maximum query sequence length in the batch.
|
||||
max_seqlen_k: int
|
||||
Maximum key sequence length in the batch.
|
||||
atten_scale: float
|
||||
The scaling of QK^T before applying softmax. Default to 1 / sqrt(headdim).
|
||||
is_causal: bool
|
||||
Whether to apply causal attention mask (e.g., for auto-regressive modeling).
|
||||
sqrt_alibi: bool
|
||||
Whether to apply abilimode
|
||||
out: (total, nheads, headdim) torch.float16, torch.bfloat16
|
||||
Returns:
|
||||
out: (total, nheads, headdim) torch.float16, torch.bfloat16
|
||||
if not q.size(-1) % 32 == 0: out shape is (total_q, nheads, q.size(-1) + (32 - q.size(-1) % 32))
|
||||
"""
|
||||
if atten_scale is None:
|
||||
atten_scale = 1.0 / (q.size(-1) ** 0.5)
|
||||
|
||||
# 判断是否pad
|
||||
cur_head = q.size(-1)
|
||||
cur_head32 = cur_head
|
||||
if not cur_head % 32 == 0:
|
||||
cur_head32 = cur_head + (32 - cur_head % 32)
|
||||
q_infer = torch.nn.functional.pad(q, [0, cur_head32 - cur_head, 0, 0], value=0)
|
||||
k_infer = torch.nn.functional.pad(k, [0, cur_head32 - cur_head, 0, 0], value=0)
|
||||
v_infer = torch.nn.functional.pad(v, [0, cur_head32 - cur_head, 0, 0], value=0)
|
||||
else:
|
||||
q_infer = q
|
||||
k_infer = k
|
||||
v_infer = v
|
||||
|
||||
if out is None:
|
||||
out = torch.empty_like(q_infer)
|
||||
# ixinfer 新接口版
|
||||
ops.infer.ixinfer_flash_attn_unpad(
|
||||
q_infer,
|
||||
k_infer,
|
||||
v_infer,
|
||||
out,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
is_causal,
|
||||
False, # need_lse =False
|
||||
atten_scale,
|
||||
sqrt_alibi,
|
||||
alibi_slopes,
|
||||
)
|
||||
if not cur_head % 32 == 0:
|
||||
out = out[:, :, :cur_head]
|
||||
return out
|
||||
|
||||
|
||||
def ref_ixinfer_flash_attn_pad(
|
||||
# [ batch num_heads seq_q head_size]
|
||||
q: torch.Tensor,
|
||||
# [ batch num_heads_k max_seq_kv head_size]
|
||||
k: torch.Tensor,
|
||||
# [ batch num_heads_k max_seq_kv head_size]
|
||||
v: torch.Tensor,
|
||||
# [ batch num_heads seq_q seq_kv] seq_kv<=max_seq_kv
|
||||
mask: torch.Tensor,
|
||||
# [ batch num_heads seq_q head_size]
|
||||
atten_scale: float = None,
|
||||
kv_seq_start: int = None,
|
||||
kv_seq_end: int = None,
|
||||
):
|
||||
head_dim = q.size(-1)
|
||||
k_effective = k[:, :, kv_seq_start:kv_seq_end, :]
|
||||
v_effective = v[:, :, kv_seq_start:kv_seq_end, :]
|
||||
# 2. q*kt softmax
|
||||
scores_qk = (
|
||||
torch.matmul(q.float(), k_effective.float().transpose(-2, -1)) * atten_scale
|
||||
)
|
||||
# softmax
|
||||
# print(scores_qk.shape,mask.shape)
|
||||
if mask is not None:
|
||||
if mask.dtype == torch.int32:
|
||||
scores_qk = scores_qk + mask * (-100000)
|
||||
elif mask.dtype == torch.float32:
|
||||
scores_qk = scores_qk + mask
|
||||
else:
|
||||
print(
|
||||
f"mask dtype is not surported {mask.dtype},now surport int32 and float32"
|
||||
)
|
||||
scores_qk = torch.nn.functional.softmax(scores_qk, dim=-1)
|
||||
# 3. x = qk_scores * v
|
||||
scores_v = torch.matmul(scores_qk, v_effective.float())
|
||||
return scores_v.half()
|
||||
|
||||
|
||||
def ixinfer_flash_attn_pad(
|
||||
# [ batch num_heads seq_q head_size]
|
||||
q: torch.Tensor,
|
||||
# [ batch num_heads_k max_seq_kv head_size]
|
||||
k: torch.Tensor,
|
||||
# [ batch num_heads_k max_seq_kv head_size]
|
||||
v: torch.Tensor,
|
||||
# [ batch num_heads seq_q seq_kv] seq_kv<=max_seq_kv
|
||||
mask: torch.Tensor,
|
||||
# [ batch num_heads seq_q head_size]
|
||||
atten_scale: float = None,
|
||||
kv_seq_start: int = None,
|
||||
kv_seq_end: int = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
q: (batch_size, num_head, seq_len_q, head_dim) torch.float16, torch.bfloat16
|
||||
k: (batch_size, num_head_kv, seq_len_kv, head_dim) torch.float16, torch.bfloat16
|
||||
v: (batch_size, num_head_kv, seq_len_kv, head_dim) torch.float16, torch.bfloat16
|
||||
mask: (batch_size, num_head, seq_len_q, kv_seq_start:kv_seq_end) torch.int32, torch.int64, torch.float32
|
||||
atten_scale: float
|
||||
The scaling of QK^T before applying softmax. Default to 1 / sqrt(headdim).
|
||||
kv_seq_start: int
|
||||
kv sequence start index used for computation in the batch
|
||||
kv_seq_end: int
|
||||
kv sequence end index used for computation in the batch.
|
||||
Returns:
|
||||
out: (batch_size, num_head, seq_len_q, head_dim) torch.float16, torch.bfloat16
|
||||
"""
|
||||
|
||||
# 判断是否pad
|
||||
cur_head = q.size(-1)
|
||||
cur_head32 = cur_head
|
||||
if not cur_head % 32 == 0:
|
||||
cur_head32 = cur_head + (32 - cur_head % 32)
|
||||
q_infer = torch.nn.functional.pad(q, [0, cur_head32 - cur_head, 0, 0], value=0)
|
||||
k_infer = torch.nn.functional.pad(k, [0, cur_head32 - cur_head, 0, 0], value=0)
|
||||
v_infer = torch.nn.functional.pad(v, [0, cur_head32 - cur_head, 0, 0], value=0)
|
||||
else:
|
||||
q_infer = q
|
||||
k_infer = k
|
||||
v_infer = v
|
||||
|
||||
if atten_scale is None:
|
||||
atten_scale = 1.0 / (q.size(-1) ** 0.5)
|
||||
if kv_seq_start is None or kv_seq_end is None:
|
||||
kv_seq_start = 0
|
||||
kv_seq_end = k.size(-2) # kv seq len
|
||||
elif kv_seq_start < 0 or kv_seq_end > k.size(-2) or kv_seq_start >= kv_seq_end:
|
||||
raise NotImplementedError(
|
||||
"must kv_seq_start<0 or kv_seq_end>k.size(-2) or kv_seq_start>=kv_seq_end!"
|
||||
)
|
||||
out_shape = list(q_infer.shape)
|
||||
out = torch.empty(out_shape, dtype=q.dtype, device=q.device)
|
||||
if mask is not None:
|
||||
ops.infer.ixinfer_flash_attn_pad_fwd(
|
||||
q_infer, k_infer, v_infer, mask, out, atten_scale, kv_seq_start, kv_seq_end
|
||||
)
|
||||
else:
|
||||
ops.infer.ixinfer_flash_attn_pad_fwd_nomask(
|
||||
q_infer, k_infer, v_infer, out, atten_scale, kv_seq_start, kv_seq_end
|
||||
)
|
||||
if not cur_head % 32 == 0:
|
||||
out = out[:, :, :, :cur_head]
|
||||
return out
|
||||
350
ixformer_sdk/inference/functions/flash_attn_lib.py
Normal file
350
ixformer_sdk/inference/functions/flash_attn_lib.py
Normal file
@@ -0,0 +1,350 @@
|
||||
import math
|
||||
from typing import List, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
from .flash_attn import ixinfer_flash_attn_unpad
|
||||
|
||||
__all__ = [
|
||||
"flash_attn_varlen_func",
|
||||
"ref_flash_attn_varlen_func",
|
||||
"flash_attn_func",
|
||||
"ref_flash_attn_func",
|
||||
]
|
||||
|
||||
|
||||
def ref_flash_attn_varlen_func(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
cu_seqlens_q: torch.Tensor,
|
||||
cu_seqlens_k: torch.Tensor,
|
||||
max_seqlen_q: int,
|
||||
max_seqlen_k: int,
|
||||
dropout_p: float = 0.0,
|
||||
softmax_scale: float = None,
|
||||
causal: bool = False,
|
||||
return_attn_probs: bool = False,
|
||||
):
|
||||
if return_attn_probs:
|
||||
raise NotImplementedError("return_attn_probs not supported!")
|
||||
out = torch.zeros_like(q)
|
||||
unpad_causal_torch(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
out,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
torch.float16,
|
||||
softmax_scale,
|
||||
causal,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def unpad_causal_torch(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
max_seq_len_q,
|
||||
max_seq_len_kv,
|
||||
dtype,
|
||||
atten_scale,
|
||||
is_causal=True,
|
||||
):
|
||||
|
||||
head_num = q.size(1)
|
||||
head_num_kv = k.size(1)
|
||||
head_dim = q.size(2)
|
||||
|
||||
assert head_num % head_num_kv == 0
|
||||
if atten_scale == None:
|
||||
atten_scale = 1.0 / (q.size(-1) ** 0.5)
|
||||
# tokens,head_num,head_dim
|
||||
if head_num != head_num_kv:
|
||||
# k = k.repeat(1, head_num//head_num_kv, 1)#[0,1,2,0,1,2,0,1,2,0,1,2]
|
||||
# v = v.repeat(1, head_num//head_num_kv, 1)
|
||||
|
||||
k = repeat_kv(k, head_num // head_num_kv) # [0,0,0,0,1,1,1,1,2,2,2,2] GROUP
|
||||
v = repeat_kv(v, head_num // head_num_kv)
|
||||
|
||||
batch_size = cu_seqlens_q.size(0) - 1
|
||||
|
||||
for i in range(batch_size):
|
||||
q_start_index = cu_seqlens_q[i]
|
||||
q_end_index = cu_seqlens_q[i + 1]
|
||||
cur_q_len = q_end_index - q_start_index
|
||||
# 1*seq_len,head_num,head_dim
|
||||
cur_q = q[q_start_index:q_end_index]
|
||||
|
||||
k_start_index = cu_seqlens_k[i]
|
||||
k_end_index = cu_seqlens_k[i + 1]
|
||||
cur_k_len = k_end_index - k_start_index
|
||||
|
||||
cur_k = k[k_start_index:k_end_index]
|
||||
cur_v = v[k_start_index:k_end_index]
|
||||
|
||||
# mask = torch.tril(torch.ones([cur_q_len, cur_k_len], dtype=torch.bool)).cuda()
|
||||
# mask = mask.unsqueeze(0).unsqueeze(0)
|
||||
if is_causal:
|
||||
# Create attention mask.
|
||||
attn_mask = torch.triu(
|
||||
torch.ones(cur_q_len, cur_k_len, dtype=dtype), diagonal=1
|
||||
)
|
||||
attn_mask = attn_mask * torch.finfo(dtype).min
|
||||
attn_mask = attn_mask.to(dtype=dtype, device="cuda")
|
||||
else:
|
||||
attn_mask = None
|
||||
|
||||
ref_output = ref_masked_attention(
|
||||
cur_q,
|
||||
cur_k,
|
||||
cur_v,
|
||||
atten_scale,
|
||||
attn_mask=attn_mask,
|
||||
)
|
||||
output[q_start_index:q_end_index].copy_(ref_output)
|
||||
|
||||
|
||||
def ref_masked_attention(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
scale: float,
|
||||
attn_mask=None,
|
||||
) -> torch.Tensor:
|
||||
query = query * scale
|
||||
dtype = query.dtype
|
||||
device = query.device
|
||||
query = query.to(torch.float32).cpu()
|
||||
key = key.to(torch.float32).cpu()
|
||||
value = value.to(torch.float32).cpu()
|
||||
attn = torch.einsum("qhd,khd->hqk", query, key)
|
||||
if attn_mask is not None:
|
||||
attn_mask = attn_mask.cpu()
|
||||
attn = attn + attn_mask
|
||||
attn = torch.softmax(attn, dim=-1)
|
||||
out = torch.einsum("hqk,khd->qhd", attn, value)
|
||||
out = out.to(device).to(dtype)
|
||||
return out
|
||||
|
||||
|
||||
def flash_attn_varlen_func(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
cu_seqlens_q: torch.Tensor,
|
||||
cu_seqlens_k: torch.Tensor,
|
||||
max_seqlen_q: int,
|
||||
max_seqlen_k: int,
|
||||
dropout_p: float = 0.0,
|
||||
softmax_scale: float = None,
|
||||
causal: bool = False,
|
||||
return_attn_probs: bool = False,
|
||||
out: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
q: (total_q, nheads, headdim) torch.float16, torch.bfloat16
|
||||
where total_q = total number of query tokens in the batch.
|
||||
k: (total_k, nheads_k, headdim) torch.float16, torch.bfloat16
|
||||
where total_k = total number of key tokens in the batch.
|
||||
v: (total_k, nheads_k, headdim) torch.float16, torch.bfloat16
|
||||
cu_seqlens_q: (batch_size + 1) torch.int32
|
||||
The cumulative sequence lengths of the sequences in the batch, used to index into q.
|
||||
cu_seqlens_k: (batch_size + 1) torch.int32
|
||||
The cumulative sequence lengths of the sequences in the batch, used to index into kv.
|
||||
max_seqlen_q: int
|
||||
Maximum query sequence length in the batch.
|
||||
max_seqlen_k: int
|
||||
Maximum key sequence length in the batch.
|
||||
dropout_p: float
|
||||
Dropout probability. dropout_p should be set to 0.0 during evaluation
|
||||
softmax_scale: float
|
||||
The scaling of QK^T before applying softmax. Default to 1 / sqrt(headdim).
|
||||
causal: bool
|
||||
Whether to apply causal attention mask (e.g., for auto-regressive modeling).
|
||||
return_attn_probs: bool
|
||||
Whether to return the attention probabilities. This option is for testing only. The returned probabilities are not guaranteed to be correct (they might not have the right scaling).
|
||||
out: (total, nheads, headdim) torch.float16, torch.bfloat16
|
||||
Returns:
|
||||
out: (total, nheads, headdim) torch.float16, torch.bfloat16
|
||||
"""
|
||||
|
||||
|
||||
assert len(q.shape) == 3, "q.shape != [total_q, nheads, head_dim]"
|
||||
assert len(k.shape) == 3, "k.shape != [total_k, nheads_k, head_dim]"
|
||||
assert len(v.shape) == 3, "v.shape != [total_k, nheads_k, head_dim]"
|
||||
assert len(cu_seqlens_q.shape) == 1, "cu_seqlens_q.shape != [batch_size+1]"
|
||||
assert len(cu_seqlens_k.shape) == 1, "cu_seqlens_k.shape != [batch_size+1]"
|
||||
|
||||
if return_attn_probs:
|
||||
raise NotImplementedError("return_attn_probs not supported!")
|
||||
atten_scale = softmax_scale
|
||||
training = q.requires_grad
|
||||
nheads = q.size(1)
|
||||
nheads_k = k.size(1)
|
||||
if training:
|
||||
raise NotImplementedError("not support training!")
|
||||
else: # 推理支持group query attention
|
||||
assert nheads % nheads_k == 0
|
||||
return ixinfer_flash_attn_unpad(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
causal,
|
||||
atten_scale,
|
||||
out=out,
|
||||
)
|
||||
|
||||
|
||||
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
|
||||
"""torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
|
||||
if len(x.shape) == 4:
|
||||
batch, seq_len, n_kv_heads, head_dim = x.shape
|
||||
elif len(x.shape) == 3:
|
||||
tokens, n_kv_heads, head_dim = x.shape
|
||||
if n_rep == 1:
|
||||
return x
|
||||
if len(x.shape) == 4:
|
||||
return (
|
||||
x[:, :, :, None, :]
|
||||
.expand(batch, seq_len, n_kv_heads, n_rep, head_dim)
|
||||
.reshape(batch, seq_len, n_kv_heads * n_rep, head_dim)
|
||||
)
|
||||
elif len(x.shape) == 3:
|
||||
return (
|
||||
x[:, :, None, :]
|
||||
.expand(tokens, n_kv_heads, n_rep, head_dim)
|
||||
.reshape(tokens, n_kv_heads * n_rep, head_dim)
|
||||
)
|
||||
|
||||
|
||||
def mha(q, k, v, atten_scale, is_causal):
|
||||
q = q.permute(0, 2, 1, 3).contiguous() # batch num_head seq_len head_dim
|
||||
k = k.permute(0, 2, 1, 3).contiguous()
|
||||
v = v.permute(0, 2, 1, 3).contiguous()
|
||||
|
||||
# 2. q*kt softmax
|
||||
scores_qk = torch.matmul(q.float(), k.float().transpose(-2, -1)) * atten_scale
|
||||
q_seq_len = q.size(2)
|
||||
kv_seq_len = k.size(2)
|
||||
if is_causal:
|
||||
# Create attention mask.
|
||||
attn_mask = torch.triu(
|
||||
torch.ones(q_seq_len, kv_seq_len, dtype=torch.int), diagonal=1
|
||||
)
|
||||
attn_mask = attn_mask.to(dtype=torch.int, device="cuda")
|
||||
else:
|
||||
attn_mask = None
|
||||
# softmax
|
||||
# print(scores_qk.shape,attn_mask.shape)
|
||||
if attn_mask is not None:
|
||||
# print(scores_qk.shape,attn_mask.shape)
|
||||
scores_qk = scores_qk + attn_mask * (-100000)
|
||||
scores_qk = torch.nn.functional.softmax(scores_qk, dim=-1)
|
||||
# 3. x = qk_scores * v
|
||||
scores_v = torch.matmul(scores_qk, v.float())
|
||||
scores_v = scores_v.half()
|
||||
return scores_v.permute(0, 2, 1, 3).contiguous()
|
||||
|
||||
|
||||
def ref_flash_attn_func(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
dropout_p: float = 0.0,
|
||||
softmax_scale: float = None,
|
||||
causal: bool = False,
|
||||
return_attn_probs: bool = False,
|
||||
):
|
||||
if return_attn_probs:
|
||||
raise NotImplementedError("return_attn_probs not supported!")
|
||||
head_num = q.size(2)
|
||||
head_num_kv = k.size(2)
|
||||
if head_num != head_num_kv:
|
||||
k = repeat_kv(k, head_num // head_num_kv) # [0,0,0,0,1,1,1,1,2,2,2,2] GROUP
|
||||
v = repeat_kv(v, head_num // head_num_kv)
|
||||
output_pt = mha(q, k, v, softmax_scale, causal)
|
||||
return output_pt
|
||||
|
||||
|
||||
def flash_attn_func(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
dropout_p: float = 0.0,
|
||||
softmax_scale: float = None,
|
||||
causal: bool = False,
|
||||
return_attn_probs: bool = False,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
q: (batch_size, seqlen, nheads, headdim) torch.float16, torch.bfloat16
|
||||
k: (batch_size, seqlen, nheads_k, headdim) torch.float16, torch.bfloat16
|
||||
v: (batch_size, seqlen, nheads_k, headdim) torch.float16, torch.bfloat16
|
||||
dropout_p: float
|
||||
Dropout probability. dropout_p should be set to 0.0 during evaluation
|
||||
softmax_scale: float
|
||||
The scaling of QK^T before applying softmax. Default to 1 / sqrt(headdim).
|
||||
causal: bool
|
||||
Whether to apply causal attention mask (e.g., for auto-regressive modeling).
|
||||
return_attn_probs: bool
|
||||
Whether to return the attention probabilities. This option is for testing only. The returned probabilities are not guaranteed to be correct (they might not have the right scaling).
|
||||
Returns:
|
||||
Tensor: (total, nheads, headdim) torch.float16, torch.bfloat16
|
||||
"""
|
||||
|
||||
if return_attn_probs:
|
||||
raise NotImplementedError("return_attn_probs not supported!")
|
||||
atten_scale = softmax_scale
|
||||
training = q.requires_grad
|
||||
|
||||
q_dim = q.dim()
|
||||
assert q_dim == 4
|
||||
|
||||
batch_size, max_seqlen_q, nheads, head_dim = q.shape
|
||||
_, max_seqlen_k, nheads_k, head_dim_k = k.shape
|
||||
assert head_dim == head_dim_k
|
||||
if training:
|
||||
raise NotImplementedError("not support training!")
|
||||
else: # 推理支持group query attention
|
||||
assert nheads % nheads_k == 0
|
||||
|
||||
q = q.view(batch_size * max_seqlen_q, nheads, head_dim)
|
||||
k = k.view(batch_size * max_seqlen_k, nheads_k, head_dim)
|
||||
v = v.view(batch_size * max_seqlen_k, nheads_k, head_dim)
|
||||
|
||||
cu_seqlens_q = torch.ones([batch_size + 1]) * max_seqlen_q
|
||||
cu_seqlens_q[0] = 0
|
||||
cu_seqlens_k = torch.ones([batch_size + 1]) * max_seqlen_k
|
||||
cu_seqlens_k[0] = 0
|
||||
cu_seqlens_q = cu_seqlens_q.cuda().int()
|
||||
cu_seqlens_k = cu_seqlens_k.cuda().int()
|
||||
cu_seqlens_q = torch.cumsum(cu_seqlens_q, dim=0).int()
|
||||
cu_seqlens_k = torch.cumsum(cu_seqlens_k, dim=0).int()
|
||||
output = ixinfer_flash_attn_unpad(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
max_seqlen_q,
|
||||
max_seqlen_k,
|
||||
causal,
|
||||
atten_scale,
|
||||
sqrt_alibi=False,
|
||||
alibi_slopes=None,
|
||||
)
|
||||
return output.view(batch_size, max_seqlen_q, nheads, head_dim)
|
||||
76
ixformer_sdk/inference/functions/fused_rope.py
Normal file
76
ixformer_sdk/inference/functions/fused_rope.py
Normal file
@@ -0,0 +1,76 @@
|
||||
from typing import List, Tuple, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
# adding by xuelu.peng 20240417
|
||||
# from https://github.com/NVIDIA/apex/blob/master/apex/transformer/functional/fused_rope.py#L59
|
||||
__all__ = ["fused_apply_rotary_pos_emb", "ref_fused_apply_rotary_pos_emb"]
|
||||
|
||||
# Copied from Megatron-Core for testing.
|
||||
# https://github.com/NVIDIA/Megatron-LM/blob/5f2877d85cb26e47ce6dcdae4b80adf376abf4e8/megatron/core/models/common/embeddings/rotary_pos_embedding.py#L139
|
||||
def apply_rotary_pos_emb(t: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
|
||||
"""Apply rotary positional embedding to input tensor T.
|
||||
|
||||
check https://kexue.fm/archives/8265 for detailed formulas
|
||||
|
||||
Arguments:
|
||||
t (Tensor): Input tensor T is of shape [seq_length, ... , dim]
|
||||
freqs (Tensor): Rotary Positional embedding tensor freq is of shape [seq_length, ..., dim]
|
||||
|
||||
Returns:
|
||||
Tensor: The input tensor after applying RoPE
|
||||
"""
|
||||
rot_dim = freqs.shape[-1]
|
||||
|
||||
# ideally t_pass is empty so rotary pos embedding is applied to all tensor t
|
||||
t, t_pass = t[..., :rot_dim], t[..., rot_dim:]
|
||||
|
||||
# first part is cosine component
|
||||
# second part is sine component, need to change signs with _rotate_half method
|
||||
cos_ = torch.cos(freqs).to(t.dtype)
|
||||
sin_ = torch.sin(freqs).to(t.dtype)
|
||||
|
||||
t = (t * cos_) + (_rotate_half(t) * sin_)
|
||||
return torch.cat((t, t_pass), dim=-1)
|
||||
|
||||
|
||||
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
|
||||
"""Change sign so the last dimension becomes [-odd, +even]
|
||||
|
||||
Arguments:
|
||||
x (Tensor): Input tensor
|
||||
|
||||
Returns:
|
||||
Tensor: Tensor rotated half
|
||||
"""
|
||||
|
||||
x1, x2 = torch.chunk(x, 2, dim=-1)
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def ref_fused_apply_rotary_pos_emb(
|
||||
t: torch.Tensor, freqs: torch.Tensor, transpose_output_memory: bool = False
|
||||
):
|
||||
output_unfused = apply_rotary_pos_emb(t, freqs)
|
||||
return output_unfused
|
||||
|
||||
|
||||
def fused_apply_rotary_pos_emb(
|
||||
t: torch.Tensor,
|
||||
freqs: torch.Tensor,
|
||||
transpose_output_memory: bool = False,
|
||||
) -> torch.Tensor:
|
||||
|
||||
"""
|
||||
Args:
|
||||
t: (sequence length,batch size,head num,head_dim) torch.float16, torch.bfloat16, torch.float32
|
||||
freqs: (sequence length,1 ,1, head_dim) torch.float32
|
||||
transpose_output_memory: bool
|
||||
Default to False. Whether to transpose the 's' and 'b' dimension of the output's underlying memory format. This is very helpful when you want to get a contiguous tensor after calling `output.transpose(0, 1)`.
|
||||
|
||||
Returns:
|
||||
Tensor: (sequence length,batch size,head num,head_dim) torch.float16, torch.bfloat16, torch.float32
|
||||
"""
|
||||
output = ops.train.fused_rope_forward(t, freqs, transpose_output_memory)
|
||||
return output
|
||||
48
ixformer_sdk/inference/functions/gemv.py
Normal file
48
ixformer_sdk/inference/functions/gemv.py
Normal file
@@ -0,0 +1,48 @@
|
||||
import os
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = ["gemv", "ref_gemv"]
|
||||
|
||||
|
||||
def ref_gemv(x: torch.Tensor, A: torch.Tensor, gemv_max_batch: int = 1):
|
||||
output = torch.nn.functional.linear(x, A)
|
||||
return output
|
||||
|
||||
|
||||
def gemv_conditions(input, weight, gemv_max_batch):
|
||||
# gemv 使用的条件 input:[m,k] weight:[n,k]
|
||||
# 1. m<=gemv_max_batch
|
||||
# 2. k%2==0 n%2==0
|
||||
# 3. bias is None
|
||||
input = input.view(-1, input.shape[-1])
|
||||
weight = weight.view(-1, weight.shape[-1])
|
||||
m = input.shape[0]
|
||||
k = input.shape[1]
|
||||
n = weight.shape[0]
|
||||
if m <= gemv_max_batch and k % 2 == 0 and n % 2 == 0:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def gemv(x: torch.Tensor, A: torch.Tensor, gemv_max_batch: int = 1):
|
||||
|
||||
"""
|
||||
Args:
|
||||
x: (..., k) torch.float16, torch.bfloat16
|
||||
A: (n,k) torch.float16, torch.bfloat16, torch.float32
|
||||
gemv_max_batch: int
|
||||
用于是否满足gemv使用条件的判断,目前只支持到1
|
||||
Returns:
|
||||
Tensor: (..., n) torch.float16, torch.bfloat16
|
||||
"""
|
||||
disable_infer_gemm_ex = os.getenv("DISABLE_INFER_GEMM_EX", "0")
|
||||
use_gemv = gemv_conditions(x, A, gemv_max_batch) and disable_infer_gemm_ex != "1"
|
||||
assert use_gemv == True
|
||||
output_shape = list(x.shape)
|
||||
output_shape[-1] = A.shape[0]
|
||||
output = x.new_empty(output_shape)
|
||||
output = ops.infer.linear_ex(x, A, None, output)
|
||||
return output
|
||||
120
ixformer_sdk/inference/functions/groupnorm.py
Normal file
120
ixformer_sdk/inference/functions/groupnorm.py
Normal file
@@ -0,0 +1,120 @@
|
||||
from typing import List, Tuple, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
from torch.autograd.function import Function, FunctionCtx
|
||||
|
||||
import ixformer
|
||||
|
||||
__all__ = [
|
||||
"group_norm",
|
||||
"ref_group_norm",
|
||||
"ref_fused_group_norm_silu",
|
||||
"fused_group_norm_silu",
|
||||
"ref_fused_group_norm_silu_nhwc",
|
||||
"fused_group_norm_silu_nhwc"
|
||||
]
|
||||
def is_channels_last(ten):
|
||||
return torch._prims_common.suggest_memory_format(ten) == torch.channels_last
|
||||
|
||||
def ref_group_norm(input, num_groups, weight, bias, eps):
|
||||
output = torch.nn.functional.group_norm(input, num_groups, weight, bias, eps)
|
||||
return output
|
||||
|
||||
#group_norm官方接口，如果input是nhwc(channel_last),输出则不是channel_last,而是nchw；如果input是nchw，那么输出也是nchw
|
||||
def group_norm(
|
||||
input: torch.Tensor,
|
||||
num_groups: int,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
eps: float = 1e-05,
|
||||
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (n,c,h,w) or (n,c,h) or (n,h,w,c) torch.float16
|
||||
"contiguous_format":(n,c,h,w) or (n,c,h) "channels_last": (n,h,w,c)
|
||||
num_groups: int
|
||||
weight: (c) torch.float16
|
||||
bias: (c) torch.float16
|
||||
eps: float
|
||||
Returns:
|
||||
Tensor: (n,c,h,w) torch.float16
|
||||
"""
|
||||
|
||||
is_nhwc=is_channels_last(input)
|
||||
out = ops.infer.groupnorm(input, num_groups, weight, bias, eps, is_nhwc, 0)
|
||||
if is_nhwc:
|
||||
out=out.permute(0,3,1,2).contiguous()
|
||||
return out
|
||||
def ref_fused_group_norm_silu_nhwc(
|
||||
input: torch.Tensor,
|
||||
num_groups: int,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
eps: float = 1e-05,
|
||||
act_type:int = 0
|
||||
):
|
||||
output = torch.nn.functional.group_norm(input.permute(0,3,1,2).contiguous(), num_groups, weight, bias, eps)
|
||||
if act_type:
|
||||
output = output * torch.sigmoid(output)
|
||||
output = output.permute(0,2,3,1).contiguous()
|
||||
return output
|
||||
#为了减少permute/contiguous,新接口支持输入输出都是nhwc的，融合silu
|
||||
def fused_group_norm_silu_nhwc(
|
||||
input: torch.Tensor,
|
||||
num_groups: int,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
eps: float = 1e-05,
|
||||
act_type:int = 0
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (n,h,w,c) torch.float16
|
||||
num_groups: int
|
||||
weight: (c) torch.float16
|
||||
bias: (c) torch.float16
|
||||
eps: float
|
||||
act_type: int
|
||||
0 or 1,if act_type=1, silu
|
||||
Returns:
|
||||
Tensor: (n,h,w,c) torch.float16
|
||||
|
||||
"""
|
||||
out = ops.infer.groupnorm(input, num_groups, weight, bias, eps, True, act_type)
|
||||
return out
|
||||
|
||||
def ref_fused_group_norm_silu(
|
||||
input: torch.Tensor,
|
||||
num_groups: int,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
eps: float = 1e-05,
|
||||
):
|
||||
output = torch.nn.functional.group_norm(input, num_groups, weight, bias, eps)
|
||||
output = output * torch.sigmoid(output)
|
||||
return output
|
||||
|
||||
|
||||
def fused_group_norm_silu(
|
||||
input: torch.Tensor,
|
||||
num_groups: int,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
eps: float = 1e-05,
|
||||
):
|
||||
|
||||
"""
|
||||
Args:
|
||||
input: (n,c,h,w) or (n,c,h) torch.float16
|
||||
num_groups: int
|
||||
weight: (c) torch.float16
|
||||
bias: (c) torch.float16
|
||||
eps: float
|
||||
Returns:
|
||||
output: (n,c,h,w) or (n,c,h) torch.float16
|
||||
"""
|
||||
|
||||
out = ops.infer.groupnorm(input, num_groups, weight, bias, eps, False, 1)
|
||||
return out
|
||||
32
ixformer_sdk/inference/functions/i8w8o32.py
Normal file
32
ixformer_sdk/inference/functions/i8w8o32.py
Normal file
@@ -0,0 +1,32 @@
|
||||
import os
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = ["i8w8o32", "ref_i8w8o32"]
|
||||
|
||||
|
||||
def ref_i8w8o32(input: torch.Tensor, weight: torch.Tensor):
|
||||
output = torch.nn.functional.linear(input.float(), weight.float()).int()
|
||||
return output
|
||||
|
||||
|
||||
def i8w8o32(input: torch.Tensor, weight: torch.Tensor):
|
||||
|
||||
"""
|
||||
Args:
|
||||
input: (bs, ic) torch.int8
|
||||
weight: (oc, ic) torch.int8
|
||||
Returns:
|
||||
Tensor: (bs, oc)) torch.int32
|
||||
"""
|
||||
if not torch.is_tensor(input):
|
||||
raise RuntimeError("Not impl.")
|
||||
output_shape = list(input.shape)
|
||||
output_shape[-1] = weight.size(0)
|
||||
output = torch.empty(output_shape, dtype=torch.int32, device=input.device)
|
||||
ic_dim = input.size(-1)
|
||||
input = input.view(-1, ic_dim)
|
||||
ops.infer.linear_i8w8o32(input.view(-1, ic_dim), weight, output)
|
||||
return output
|
||||
429
ixformer_sdk/inference/functions/layernorm.py
Normal file
429
ixformer_sdk/inference/functions/layernorm.py
Normal file
@@ -0,0 +1,429 @@
|
||||
from typing import List, Tuple, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"layer_norm",
|
||||
"ref_layer_norm",
|
||||
"residual_layer_norm",
|
||||
"ref_residual_layer_norm",
|
||||
"ref_residual_layer_norm_bias_alpha",
|
||||
"residual_layer_norm_bias_alpha",
|
||||
"ref_layer_norm_2sb_fused",
|
||||
"layer_norm_2sb_fused",
|
||||
]
|
||||
|
||||
|
||||
def ref_layer_norm(
|
||||
input: torch.Tensor,
|
||||
normalized_shape: List[int],
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
eps: float = 1e-5,
|
||||
output: torch.Tensor = None,
|
||||
):
|
||||
|
||||
if weight is None or bias is None or weight.dim() > 1 or bias.dim() > 1:
|
||||
raise NotImplementedError(
|
||||
"layer_norm only support weight.dim() ==1 and bias.dim()==1!"
|
||||
)
|
||||
if normalized_shape == None:
|
||||
norm_size = weight.size(-1)
|
||||
normalized_shape = [norm_size]
|
||||
else:
|
||||
if (
|
||||
isinstance(normalized_shape, list) or isinstance(normalized_shape, tuple)
|
||||
) and len(normalized_shape) == 1:
|
||||
norm_size = normalized_shape[0]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"layer_norm(): argument 'normalized_shape' (position 2) must be tuple of ints and length of tuple is equal to 1, not {type(normalized_shape)}"
|
||||
)
|
||||
if norm_size != weight.size(-1):
|
||||
raise ValueError(f"layer_norm(): argument 'norm_size' must == weight.size(-1)")
|
||||
|
||||
norm_out = torch.nn.functional.layer_norm(
|
||||
input, normalized_shape, weight, bias, eps=eps
|
||||
)
|
||||
if output is not None:
|
||||
assert output.shape == norm_out.shape
|
||||
output.copy_(norm_out)
|
||||
else:
|
||||
output = norm_out
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def layer_norm(
|
||||
input: torch.Tensor,
|
||||
normalized_shape: List[int],
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
eps: float = 1e-5,
|
||||
output: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
This function is deprecated, please use residual_layer_norm.
|
||||
等价实现：
|
||||
torch.nn.functional.layer_norm( input, normalized_shape, weight, bias, eps=0.000001)
|
||||
Args:
|
||||
input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
normalized_shape: list[int]
|
||||
weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
eps: float32
|
||||
Returns:
|
||||
output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
"""
|
||||
if weight is None or bias is None or weight.dim() > 1 or bias.dim() > 1:
|
||||
raise NotImplementedError(
|
||||
"layer_norm only support weight.dim() ==1 and bias.dim()==1!"
|
||||
)
|
||||
if normalized_shape == None:
|
||||
norm_size = weight.size(-1)
|
||||
normalized_shape = [norm_size]
|
||||
else:
|
||||
if (
|
||||
isinstance(normalized_shape, list) or isinstance(normalized_shape, tuple)
|
||||
) and len(normalized_shape) == 1:
|
||||
norm_size = normalized_shape[0]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"layer_norm(): argument 'normalized_shape' (position 2) must be tuple of ints and length of tuple is equal to 1, not {type(normalized_shape)}"
|
||||
)
|
||||
if norm_size != weight.size(-1):
|
||||
raise ValueError(f"layer_norm(): argument 'norm_size' must == weight.size(-1)")
|
||||
if output is None:
|
||||
output = torch.empty_like(input)
|
||||
ops.infer.layer_norm(input, weight, bias, None, output, eps)
|
||||
return output
|
||||
|
||||
|
||||
def ref_residual_layer_norm(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
residual: torch.Tensor = None,
|
||||
residual_bias: torch.Tensor = None,
|
||||
eps: float = 1e-5,
|
||||
output: torch.Tensor = None,
|
||||
residual_output: torch.Tensor = None,
|
||||
):
|
||||
normalized_shape = [weight.size(-1)]
|
||||
|
||||
if residual_bias is not None:
|
||||
input = input + residual_bias
|
||||
|
||||
if residual is not None:
|
||||
residual_output = torch.add(input, residual, out=residual_output)
|
||||
input = residual_output
|
||||
|
||||
norm_out = torch.nn.functional.layer_norm(
|
||||
input, normalized_shape, weight, bias, eps=eps
|
||||
)
|
||||
|
||||
if output is None:
|
||||
output = norm_out
|
||||
else:
|
||||
output.copy_(norm_out)
|
||||
|
||||
return output, residual_output
|
||||
|
||||
|
||||
def residual_layer_norm(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
residual: torch.Tensor = None,
|
||||
residual_bias: torch.Tensor = None,
|
||||
eps: float = 1e-5,
|
||||
output: torch.Tensor = None,
|
||||
residual_output: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
residual: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
residual_bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
eps: float32
|
||||
Returns:
|
||||
output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 If set to None, an inplace operation will be performed on input.
|
||||
residual_output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 If set to None, an inplace operation will be performed on residual.
|
||||
"""
|
||||
if residual is None:
|
||||
if output is None:
|
||||
output = torch.empty(input.shape, device=input.device, dtype=input.dtype)
|
||||
ops.infer.layer_norm(input, weight, bias, residual_bias, output, eps)
|
||||
else:
|
||||
ops.infer.residual_layer_norm(
|
||||
input,
|
||||
residual,
|
||||
weight,
|
||||
bias,
|
||||
residual_bias,
|
||||
output,
|
||||
residual_output,
|
||||
1.0,
|
||||
eps,
|
||||
False,
|
||||
)
|
||||
residual_output = residual_output if residual_output is not None else residual
|
||||
output = output if output is not None else input
|
||||
|
||||
return output, residual_output
|
||||
|
||||
|
||||
def ref_residual_layer_norm_bias_alpha(
|
||||
input: torch.Tensor,
|
||||
normalized_shape: List[int],
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
residual_bias: torch.Tensor = None,
|
||||
alpha: float = 1.0,
|
||||
eps: float = 1e-5,
|
||||
is_post_ln=False,
|
||||
):
|
||||
if (
|
||||
weight is None
|
||||
or bias is None
|
||||
or residual is None
|
||||
or weight.dim() > 1
|
||||
or bias.dim() > 1
|
||||
):
|
||||
raise NotImplementedError(
|
||||
"residual_layer_norm only support weight.dim() ==1 and bias.dim()==1!"
|
||||
)
|
||||
if normalized_shape == None:
|
||||
norm_size = weight.size(-1)
|
||||
normalized_shape = [norm_size]
|
||||
else:
|
||||
if (
|
||||
isinstance(normalized_shape, list) or isinstance(normalized_shape, tuple)
|
||||
) and len(normalized_shape) == 1:
|
||||
norm_size = normalized_shape[0]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"residual_layer_norm(): argument 'normalized_shape' (position 2) must be tuple of ints and length of tuple is equal to 1, not {type(normalized_shape)}"
|
||||
)
|
||||
|
||||
if norm_size != weight.size(-1):
|
||||
raise ValueError(
|
||||
f"residual_layer_norm(): argument 'norm_size' must == weight.size(-1)"
|
||||
)
|
||||
dtype = input.dtype
|
||||
if residual_bias is None:
|
||||
x = input.float() + residual.float() * alpha
|
||||
else:
|
||||
x = input.float() + residual.float() * alpha + residual_bias.float()
|
||||
|
||||
y = torch.nn.functional.layer_norm(
|
||||
x.to(dtype), normalized_shape, weight, bias, eps=eps
|
||||
)
|
||||
|
||||
if is_post_ln:
|
||||
return y, y
|
||||
else:
|
||||
return y, x.to(dtype)
|
||||
|
||||
|
||||
def residual_layer_norm_bias_alpha(
|
||||
input: torch.Tensor,
|
||||
normalized_shape: List[int],
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
residual_bias: torch.Tensor = None,
|
||||
alpha: float = 1.0,
|
||||
eps: float = 1e-5,
|
||||
is_post_ln=False,
|
||||
):
|
||||
"""
|
||||
等价实现：
|
||||
residual = input + residual.float() * alpha + residual_bias
|
||||
output = torch.nn.functional.layer_norm(
|
||||
residual, normalized_shape, weight, bias, eps=eps
|
||||
)
|
||||
residual = output if is_post_ln else residual
|
||||
|
||||
Args:
|
||||
input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
normalized_shape list[int]
|
||||
weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
residual: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
residual_bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
alpha: float32
|
||||
eps: float32
|
||||
is_post_ln: bool
|
||||
Returns:
|
||||
input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
residual: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
Inplace operation will be performed on residual and input.
|
||||
"""
|
||||
if (
|
||||
weight is None
|
||||
or bias is None
|
||||
or residual is None
|
||||
or weight.dim() > 1
|
||||
or bias.dim() > 1
|
||||
):
|
||||
raise NotImplementedError(
|
||||
"residual_layer_norm only support weight.dim() ==1 and bias.dim()==1!"
|
||||
)
|
||||
if normalized_shape == None:
|
||||
norm_size = weight.size(-1)
|
||||
normalized_shape = [norm_size]
|
||||
else:
|
||||
if (
|
||||
isinstance(normalized_shape, list) or isinstance(normalized_shape, tuple)
|
||||
) and len(normalized_shape) == 1:
|
||||
norm_size = normalized_shape[0]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"residual_layer_norm(): argument 'normalized_shape' (position 2) must be tuple of ints and length of tuple is equal to 1, not {type(normalized_shape)}"
|
||||
)
|
||||
|
||||
if norm_size != weight.size(-1):
|
||||
raise ValueError(
|
||||
f"residual_layer_norm(): argument 'norm_size' must == weight.size(-1)"
|
||||
)
|
||||
|
||||
ops.infer.residual_layer_norm(
|
||||
input, residual, weight, bias, residual_bias, None, None, alpha, eps, is_post_ln
|
||||
)
|
||||
|
||||
return input, residual
|
||||
|
||||
|
||||
def ref_layer_norm_2sb_fused(
|
||||
input: torch.Tensor,
|
||||
normalized_shape: List[int],
|
||||
weight1: torch.Tensor,
|
||||
bias1: torch.Tensor,
|
||||
weight2: torch.Tensor,
|
||||
bias2: torch.Tensor,
|
||||
eps: float = 1e-5,
|
||||
):
|
||||
assert input.shape[-1] <= 16384
|
||||
if not (input.dtype == torch.float16 or input.dtype == torch.bfloat16):
|
||||
raise NotImplementedError(
|
||||
"layer_norm_2sb() only support data format of float16 or bfloat16 now!"
|
||||
)
|
||||
if (
|
||||
weight1 is None
|
||||
or bias1 is None
|
||||
or weight2 is None
|
||||
or bias2 is None
|
||||
or weight1.dim() > 1
|
||||
or bias1.dim() > 1
|
||||
or weight2.dim() > 1
|
||||
or bias2.dim() > 1
|
||||
):
|
||||
raise NotImplementedError(
|
||||
"layer_norm_2sb only support weight1.dim() ==1, bias1.dim()==1, weight2.dim() ==1 and bias2.dim()==1 !"
|
||||
)
|
||||
if normalized_shape == None:
|
||||
norm_size = weight1.size(-1)
|
||||
normalized_shape = [norm_size]
|
||||
else:
|
||||
if (
|
||||
isinstance(normalized_shape, list) or isinstance(normalized_shape, tuple)
|
||||
) and len(normalized_shape) == 1:
|
||||
norm_size = normalized_shape[0]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"layer_norm_2sb(): argument 'normalized_shape' (position 2) must be tuple of ints and length of tuple is equal to 1, not {type(normalized_shape)}"
|
||||
)
|
||||
|
||||
if norm_size != weight1.size(-1) or norm_size != weight2.size(-1):
|
||||
raise ValueError(
|
||||
f"layer_norm_2sb(): argument 'norm_size' must == weight.size(-1)"
|
||||
)
|
||||
|
||||
output1 = torch.nn.functional.layer_norm(
|
||||
input, normalized_shape, weight1, bias1, eps=eps
|
||||
)
|
||||
|
||||
output2 = torch.nn.functional.layer_norm(
|
||||
input, normalized_shape, weight2, bias2, eps=eps
|
||||
)
|
||||
|
||||
return output1, output2
|
||||
|
||||
|
||||
def layer_norm_2sb_fused(
|
||||
input: torch.Tensor,
|
||||
normalized_shape: List[int],
|
||||
weight1: torch.Tensor,
|
||||
bias1: torch.Tensor,
|
||||
weight2: torch.Tensor,
|
||||
bias2: torch.Tensor,
|
||||
eps: float = 1e-5,
|
||||
):
|
||||
"""
|
||||
等价实现：
|
||||
output1 = torch.nn.functional.layer_norm(
|
||||
input, normalized_shape, weight1, bias1, eps=eps
|
||||
)
|
||||
output2 = torch.nn.functional.layer_norm(
|
||||
input, normalized_shape, weight2, bias2, eps=eps
|
||||
)
|
||||
|
||||
Args:
|
||||
input: (..., hidden_size) torch.float16, torch.bfloat16
|
||||
normalized_shape list[int]
|
||||
weight1: (hidden_size) torch.float16, torch.bfloat16
|
||||
bias1: (hidden_size) torch.float16, torch.bfloat16
|
||||
weight2: (hidden_size) torch.float16, torch.bfloat16
|
||||
bias2: (hidden_size) torch.float16, torch.bfloat16
|
||||
eps: float32
|
||||
Returns:
|
||||
output1: (..., hidden_size) torch.float16, torch.bfloat16
|
||||
output2: (..., hidden_size) torch.float16, torch.bfloat16
|
||||
"""
|
||||
assert input.shape[-1] <= 16384
|
||||
if not (input.dtype == torch.float16 or input.dtype == torch.bfloat16):
|
||||
raise NotImplementedError(
|
||||
"layer_norm_2sb() only support data format of float16 or bfloat16 now!"
|
||||
)
|
||||
if (
|
||||
weight1 is None
|
||||
or bias1 is None
|
||||
or weight2 is None
|
||||
or bias2 is None
|
||||
or weight1.dim() > 1
|
||||
or bias1.dim() > 1
|
||||
or weight2.dim() > 1
|
||||
or bias2.dim() > 1
|
||||
):
|
||||
raise NotImplementedError(
|
||||
"layer_norm_2sb only support weight1.dim() ==1, bias1.dim()==1, weight2.dim() ==1 and bias2.dim()==1 !"
|
||||
)
|
||||
if normalized_shape == None:
|
||||
norm_size = weight1.size(-1)
|
||||
normalized_shape = [norm_size]
|
||||
else:
|
||||
if (
|
||||
isinstance(normalized_shape, list) or isinstance(normalized_shape, tuple)
|
||||
) and len(normalized_shape) == 1:
|
||||
norm_size = normalized_shape[0]
|
||||
else:
|
||||
raise ValueError(
|
||||
f"layer_norm_2sb(): argument 'normalized_shape' (position 2) must be tuple of ints and length of tuple is equal to 1, not {type(normalized_shape)}"
|
||||
)
|
||||
|
||||
if norm_size != weight1.size(-1) or norm_size != weight2.size(-1):
|
||||
raise ValueError(
|
||||
f"layer_norm_2sb(): argument 'norm_size' must == weight.size(-1)"
|
||||
)
|
||||
output1 = torch.empty_like(input)
|
||||
output2 = torch.empty_like(input)
|
||||
ops.infer.layer_norm_2sb(
|
||||
input, weight1, bias1, weight2, bias2, eps, output1, output2
|
||||
)
|
||||
|
||||
return output1, output2
|
||||
277
ixformer_sdk/inference/functions/lightllm.py
Normal file
277
ixformer_sdk/inference/functions/lightllm.py
Normal file
@@ -0,0 +1,277 @@
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"lightllm_tokenattention",
|
||||
"ref_lightllm_tokenattention",
|
||||
"lightllm_destindex_copy_kv",
|
||||
"ref_lightllm_destindex_copy_kv",
|
||||
"lightllm_apply_penalty",
|
||||
"ref_lightllm_apply_penalty",
|
||||
"lightllm_glm2_rope",
|
||||
"ref_lightllm_glm2_rope",
|
||||
]
|
||||
|
||||
|
||||
def ref_lightllm_glm2_rope(
|
||||
x: torch.Tensor, # tokens,head_num,head_dim
|
||||
cos: torch.Tensor, # tokens,rotdim
|
||||
sin: torch.Tensor,
|
||||
):
|
||||
num_tokens, _, rot_dim = list(cos.shape)
|
||||
head_num = x.shape[1]
|
||||
x12 = x[:, :, : rot_dim * 2]
|
||||
x3 = x[:, :, rot_dim * 2 :]
|
||||
x12 = x12.reshape(num_tokens, head_num, rot_dim, 2)
|
||||
x1 = x12[:, :, :, 0]
|
||||
x2 = x12[:, :, :, 1]
|
||||
|
||||
# out0 = q0 * cos - q1 * sin
|
||||
# out1 = q0 * sin + q1 * cos
|
||||
# q1, q2 是沿着 head_dim维度，交叉取值的
|
||||
|
||||
q1 = x1 * cos - x2 * sin
|
||||
q2 = x2 * cos + x1 * sin
|
||||
|
||||
q12 = torch.stack([q1, q2], dim=-1)
|
||||
q12 = q12.reshape(num_tokens, head_num, -1)
|
||||
x_pytorch = torch.cat([q12, x3], dim=-1)
|
||||
return x_pytorch
|
||||
|
||||
|
||||
def lightllm_glm2_rope(
|
||||
x: torch.Tensor, # tokens,head_num,head_dim
|
||||
cos: torch.Tensor, # tokens,rotdim
|
||||
sin: torch.Tensor,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
x: (num_tokens, head_num, head_dim) torch.half
|
||||
cos: (num_tokens,1,head_dim//2//2) torch.half
|
||||
sin: (num_tokens,1,head_dim//2//2) torch.half
|
||||
Returns:
|
||||
x: (num_tokens, head_num, head_dim) torch.half
|
||||
|
||||
"""
|
||||
if isinstance(x, torch.Tensor):
|
||||
ops.infer.lightllm_glm2_rope(x, cos, sin)
|
||||
return x
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def ref_lightllm_apply_penalty(
|
||||
Logits: torch.Tensor,
|
||||
presence_penalty: torch.Tensor,
|
||||
freqency_penalty: torch.Tensor,
|
||||
p_token_ids: torch.Tensor,
|
||||
p_token_counts: torch.Tensor,
|
||||
p_cumsum_seq_len: torch.Tensor,
|
||||
p_max_len_in_batch: int,
|
||||
):
|
||||
batch_size = Logits.size(0)
|
||||
output = Logits.clone()
|
||||
for cur_batch in range(batch_size):
|
||||
cur_freqency = freqency_penalty[cur_batch]
|
||||
cur_presence = presence_penalty[cur_batch]
|
||||
cur_batch_start_index = p_cumsum_seq_len[cur_batch]
|
||||
cur_batch_end_index = p_cumsum_seq_len[cur_batch + 1]
|
||||
for token_idx in range(cur_batch_start_index, cur_batch_end_index):
|
||||
batch_ids = p_token_ids[token_idx]
|
||||
batch_ids_count = p_token_counts[token_idx]
|
||||
cur_logits = output[cur_batch][batch_ids]
|
||||
|
||||
freq_logits = cur_logits - batch_ids_count * cur_freqency
|
||||
pre_logits = freq_logits - cur_presence
|
||||
# if token_idx==0:
|
||||
# print(f"batch_ids {batch_ids} cur_logits {cur_logits} pre_logits {pre_logits}")
|
||||
output[cur_batch][batch_ids] = pre_logits
|
||||
return output
|
||||
|
||||
|
||||
def lightllm_apply_penalty(
|
||||
Logits: torch.Tensor,
|
||||
presence_penalty: torch.Tensor,
|
||||
freqency_penalty: torch.Tensor,
|
||||
p_token_ids: torch.Tensor,
|
||||
p_token_counts: torch.Tensor,
|
||||
p_cumsum_seq_len: torch.Tensor,
|
||||
p_max_len_in_batch: int,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
logits: (batch_size, vocab_size) torch.float
|
||||
presence_penalty: (batch_size) torch.float
|
||||
freqency_penalty: (batch_size) torch.float
|
||||
p_token_ids: (num_tokens) torch.int
|
||||
p_token_counts: (num_tokens) torch.int
|
||||
p_cumsum_seq_len: (batch_size+1) torch.int
|
||||
p_max_len_in_batch: int
|
||||
在一个batch中seq的最大长度
|
||||
Returns:
|
||||
logits: (batch_size, vocab_size) torch.float
|
||||
"""
|
||||
if isinstance(Logits, torch.Tensor):
|
||||
ops.infer.lightllm_apply_penalty(
|
||||
Logits,
|
||||
presence_penalty,
|
||||
freqency_penalty,
|
||||
p_token_ids,
|
||||
p_token_counts,
|
||||
p_cumsum_seq_len,
|
||||
p_max_len_in_batch,
|
||||
)
|
||||
return Logits
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def ref_lightllm_destindex_copy_kv(
|
||||
key_cache: torch.Tensor,
|
||||
mem_idx: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
):
|
||||
if key_cache.dim() != 3 or key_cache.size(-1) != 128:
|
||||
raise NotImplementedError(
|
||||
"lightllm_destindex_copy_kv only support key_cache.dim()==3 and head_size ==128 !"
|
||||
)
|
||||
output[mem_idx.long()] = key_cache
|
||||
return output
|
||||
|
||||
|
||||
def lightllm_destindex_copy_kv(
|
||||
key_cache: torch.Tensor,
|
||||
mem_idx: torch.Tensor,
|
||||
output: torch.Tensor,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
key_cache: (tokens, num_kv_heads, head_size) torch.half
|
||||
目前head_size 只支持128的情况
|
||||
mem_idx: (tokens) torch.int
|
||||
output: (max_tokens, num_kv_heads, head_size) torch.half
|
||||
Returns:
|
||||
output: (max_tokens, num_kv_heads, head_size) torch.half
|
||||
"""
|
||||
|
||||
if key_cache.dim() != 3 or key_cache.size(-1) != 128:
|
||||
raise NotImplementedError(
|
||||
"lightllm_destindex_copy_kv only support key_cache.dim()==3 and head_size ==128 !"
|
||||
)
|
||||
if isinstance(key_cache, torch.Tensor):
|
||||
ops.infer.lightllm_destindex_copy_kv(key_cache, mem_idx, output)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
return output
|
||||
|
||||
|
||||
def ref_lightllm_tokenattention(
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
reg_tokens: torch.Tensor,
|
||||
b_req_idx: torch.Tensor,
|
||||
b_seq_len: torch.Tensor,
|
||||
scale: float,
|
||||
max_context_len: int,
|
||||
):
|
||||
batch_size, tp_q_head_num_, head_dim_ = query.shape
|
||||
tp_k_head_num_ = key_cache.size(-2)
|
||||
sm_scale = scale
|
||||
curbatch_max_context_len = max_context_len
|
||||
tmp_k = torch.zeros(
|
||||
(batch_size, tp_q_head_num_, curbatch_max_context_len, head_dim_),
|
||||
dtype=query.dtype,
|
||||
device="cuda",
|
||||
)
|
||||
tmp_v = torch.zeros(
|
||||
(batch_size, tp_q_head_num_, curbatch_max_context_len, head_dim_),
|
||||
dtype=query.dtype,
|
||||
device="cuda",
|
||||
)
|
||||
mask = torch.ones([batch_size, 1, 1, curbatch_max_context_len])
|
||||
|
||||
kv_group_num = tp_q_head_num_ // tp_k_head_num_
|
||||
for cur_batch in range(batch_size):
|
||||
cur_batch_req_idx = b_req_idx[cur_batch]
|
||||
seq_len = b_seq_len[cur_batch]
|
||||
mask[cur_batch, :, :, :seq_len] = 0
|
||||
# print(f"cur_batch {cur_batch}")
|
||||
|
||||
for seq_idx in range(seq_len):
|
||||
k_loc = reg_tokens[cur_batch_req_idx][seq_idx]
|
||||
# print(k_loc)
|
||||
for cur_head in range(tp_q_head_num_):
|
||||
cur_kv_head = cur_head // kv_group_num
|
||||
tmp_k[cur_batch, cur_head, seq_idx, :] = key_cache[k_loc][cur_kv_head]
|
||||
tmp_v[cur_batch, cur_head, seq_idx, :] = value_cache[k_loc][cur_kv_head]
|
||||
mask = mask.cuda()
|
||||
# batch_size, self.tp_q_head_num_, 1, max_len_in_batch
|
||||
attn_score = (
|
||||
torch.matmul(
|
||||
query.view(batch_size, tp_q_head_num_, 1, head_dim_),
|
||||
tmp_k.transpose(-1, -2),
|
||||
)
|
||||
* sm_scale
|
||||
)
|
||||
attn_score = attn_score + mask * -1000
|
||||
attn_score = torch.softmax(attn_score, dim=-1)
|
||||
# batch_size, self.tp_q_head_num_, 1, head_dim
|
||||
py_out = torch.matmul(attn_score.to(query.dtype), tmp_v).view(
|
||||
batch_size, tp_q_head_num_, -1
|
||||
)
|
||||
return py_out
|
||||
|
||||
|
||||
def lightllm_tokenattention(
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
reg_tokens: torch.Tensor,
|
||||
b_req_idx: torch.Tensor,
|
||||
b_seq_len: torch.Tensor,
|
||||
scale: float,
|
||||
max_context_len: int,
|
||||
partition: int,
|
||||
output: torch.Tensor,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
query: (batch_size,head_num,head_dim) torch.float16, torch.bfloat16
|
||||
key_cache: (max_num_tokens, head_num_kv, head_dim) torch.float16, torch.bfloat16
|
||||
value_cache: (max_num_tokens, head_num_kv, head_dim) torch.float16, torch.bfloat16
|
||||
reg_tokens: (max_request,max_tokens) torch.int32
|
||||
目前max_tokens只支持3080
|
||||
b_req_idx: (batch_size) torch.int32
|
||||
b_req_len: (batch_size) torch.int32
|
||||
scale: float
|
||||
The scaling of QK^T before applying softmax.
|
||||
max_context_len: int
|
||||
b_seq_len.max()
|
||||
partition: int
|
||||
Returns:
|
||||
output: (batch_size,head_num,head_dim) torch.float16, torch.bfloat16
|
||||
"""
|
||||
_,max_tokens=reg_tokens.shape
|
||||
if not max_tokens == 3080:
|
||||
raise NotImplementedError(
|
||||
"lightllm_tokenattention only support reg_tokens.size(-1)==3080"
|
||||
)
|
||||
if isinstance(query, torch.Tensor):
|
||||
ops.infer.lightllm_tokenattention(
|
||||
query,
|
||||
key_cache,
|
||||
value_cache,
|
||||
reg_tokens,
|
||||
b_req_idx,
|
||||
b_seq_len,
|
||||
scale,
|
||||
max_context_len,
|
||||
partition,
|
||||
output,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
return output
|
||||
50
ixformer_sdk/inference/functions/linalg.py
Normal file
50
ixformer_sdk/inference/functions/linalg.py
Normal file
@@ -0,0 +1,50 @@
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = ["solve", "ref_slove"]
|
||||
|
||||
|
||||
def ref_slove(
|
||||
A: torch.Tensor, B: torch.Tensor, *, left: bool = True, out: torch.Tensor = None
|
||||
):
|
||||
out = torch.linalg.solve(A, B, left=left)
|
||||
return out
|
||||
|
||||
|
||||
def solve(
|
||||
A: torch.Tensor, B: torch.Tensor, *, left: bool = True, out: torch.Tensor = None
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
A: (..., n, n) torch.float
|
||||
B: (..., n) or (..., n, k) or (n,...) or (n, k) or (n) torch.float
|
||||
left: bool
|
||||
whether to solve the system AX=B or XA=B. Default: True, 目前只支持left =True
|
||||
out: (..., n, k) torch.float
|
||||
Returns:
|
||||
out: (..., n, k) torch.float
|
||||
"""
|
||||
|
||||
n = A.shape[-1]
|
||||
batch_count = A.numel() // (n * n)
|
||||
if B.dim() == 1:
|
||||
k = 1
|
||||
elif B.dim() == 2:
|
||||
if A.dim() > 2 and B.shape == (batch_count, n):
|
||||
k = 1
|
||||
else:
|
||||
nid = 0 if left else 1
|
||||
k = B.shape[nid ^ 1]
|
||||
else:
|
||||
k = B.size(B.dim() - 1 if left else B.dim() - 2)
|
||||
|
||||
if n <= 64 and k <= 64 and left:
|
||||
return ops.infer.solve(A, B, left)
|
||||
else:
|
||||
device = A.device
|
||||
cpu_A = A.cpu()
|
||||
cpu_B = B.cpu()
|
||||
cpu_res = ref_slove(A=cpu_A, B=cpu_B, left=left)
|
||||
return cpu_res.to(device)
|
||||
122
ixformer_sdk/inference/functions/linear.py
Normal file
122
ixformer_sdk/inference/functions/linear.py
Normal file
@@ -0,0 +1,122 @@
|
||||
import os
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = ["linear", "ref_linear", "mixed_type_linear", "ref_mixed_type_linear"]
|
||||
|
||||
|
||||
def ref_linear(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
act_type=-1,
|
||||
):
|
||||
output = torch.nn.functional.linear(input, weight, bias)
|
||||
if act_type == -1:
|
||||
act_fn = torch.nn.Identity()
|
||||
elif act_type == 3:
|
||||
act_fn = torch.nn.GELU()
|
||||
elif act_type == 4:
|
||||
act_fn = torch.nn.ReLU()
|
||||
elif act_type == 12:
|
||||
act_fn = torch.nn.SiLU()
|
||||
else:
|
||||
raise KeyError("act_type not supported")
|
||||
output = act_fn(output)
|
||||
return output
|
||||
|
||||
|
||||
def gemv_conditions(input, weight, bias, gemv_max_batch):
|
||||
# gemv 使用的条件 input:[m,k] weight:[n,k]
|
||||
# 1. m<=gemv_max_batch
|
||||
# 2. k%32==0 n%2==0
|
||||
# 3. bias is None
|
||||
input = input.view(-1, input.shape[-1])
|
||||
weight = weight.view(-1, weight.shape[-1])
|
||||
m = input.shape[0]
|
||||
k = input.shape[1]
|
||||
n = weight.shape[0]
|
||||
if bias is None and m <= gemv_max_batch and k % 32 == 0 and n % 2 == 0:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def linear(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
output: torch.Tensor = None,
|
||||
persistent: bool = False,
|
||||
act_type : int = -1,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (...,k) torch.float16, torch.bfloat16
|
||||
weight: (n, k) torch.float16, torch.bfloat16
|
||||
bias: (n) torch.float16, torch.bfloat16
|
||||
output: (...,n) torch.float16, torch.bfloat16
|
||||
persistent: bool
|
||||
是否限制 Gemm Kernel 的 Block 数量
|
||||
Returns:
|
||||
output: (...,n) torch.float16, torch.bfloat16
|
||||
"""
|
||||
if not input.is_contiguous():
|
||||
input = input.contiguous()
|
||||
if not weight.is_contiguous():
|
||||
weight = weight.contiguous()
|
||||
use_gemv = True
|
||||
gemv_max_batch = 1
|
||||
disable_infer_gemm_ex = os.getenv("DISABLE_INFER_GEMM_EX", "0")
|
||||
use_gemv = (
|
||||
use_gemv
|
||||
and gemv_conditions(input, weight, bias, gemv_max_batch)
|
||||
and disable_infer_gemm_ex != "1"
|
||||
)
|
||||
|
||||
if output is None:
|
||||
output_shape = list(input.shape)
|
||||
output_shape[-1] = weight.shape[0]
|
||||
output = input.new_empty(output_shape)
|
||||
|
||||
if not use_gemv:
|
||||
output = ops.infer.linear(input, weight, act_type, bias, output, persistent)
|
||||
else:
|
||||
output = ops.infer.linear_ex(input, weight, bias, output)
|
||||
return output
|
||||
|
||||
|
||||
def ref_mixed_type_linear(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
output: torch.Tensor = None,
|
||||
persistent=False, # TODO: support persistent
|
||||
):
|
||||
input = input.to(weight.dtype)
|
||||
if bias:
|
||||
bias = bias.to(weight.dtype)
|
||||
output = torch.nn.functional.linear(input, weight, bias)
|
||||
return output
|
||||
|
||||
|
||||
def mixed_type_linear(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
output: torch.Tensor = None,
|
||||
persistent=False, # TODO: support persistent
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (...,k) torch.half, torch.bfloat16
|
||||
weight: (m, k) torch.float32
|
||||
bias: not supported
|
||||
output: (...,m) torch.float32
|
||||
persistent: bool
|
||||
Returns:
|
||||
output: (...,m) torch.float32
|
||||
"""
|
||||
output = ops.infer.mixed_type_linear(input, weight, bias, output)
|
||||
return output
|
||||
212
ixformer_sdk/inference/functions/lmdeploy.py
Normal file
212
ixformer_sdk/inference/functions/lmdeploy.py
Normal file
@@ -0,0 +1,212 @@
|
||||
import math
|
||||
from typing import Literal, Optional, Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import ixformer._C._functions as CF
|
||||
import torch
|
||||
|
||||
from ixformer.core import config
|
||||
|
||||
from .linear import linear
|
||||
from .paged_attention import paged_attention as paged_attention_ixformer_impl
|
||||
|
||||
__all__ = [
|
||||
"ref_lmdeploy_paged_attention",
|
||||
"lmdeploy_paged_attention",
|
||||
]
|
||||
|
||||
weak_ref_tensor = ops.infer.weak_ref_tensor
|
||||
|
||||
|
||||
def ref_lmdeploy_paged_attention(
|
||||
output: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
num_kv_heads: torch.Tensor,
|
||||
scale: float,
|
||||
block_tables: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
block_size: int,
|
||||
max_context_len: int,
|
||||
alibi_slopes: torch.Tensor = None,
|
||||
softcap: float = 0.0,
|
||||
window_left: int = -1,
|
||||
window_right: int = -1,
|
||||
use_sqrt_alibi: bool = False,
|
||||
quant_type: int = 0,
|
||||
is_bbhh: bool = False,
|
||||
):
|
||||
assert window_right in [-1, 0]
|
||||
|
||||
if is_bbhh:
|
||||
key_cache = key_cache.permute(0, 2, 1, 3).contiguous()
|
||||
value_cache = value_cache.permute(0, 2, 1, 3).contiguous()
|
||||
|
||||
def get_alibi_mask(num_heads, seqlen, device, dtype):
|
||||
x = torch.arange(0, seqlen, device=device, dtype=torch.float32).view(-1, 1)
|
||||
y = torch.tensor(seqlen - 1, device=device, dtype=torch.float32).view(1, -1)
|
||||
offsets = -(y - x).view(1, 1, seqlen)
|
||||
return offsets
|
||||
|
||||
def ref_masked_attention(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
scale: float,
|
||||
attn_mask: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
query = query * scale
|
||||
dtype = query.dtype
|
||||
device = query.device
|
||||
query = query.to(torch.float32)
|
||||
key = key.to(torch.float32)
|
||||
value = value.to(torch.float32)
|
||||
attn = torch.einsum("qhd,khd->hqk", query, key)
|
||||
if attn_mask is not None:
|
||||
attn_mask = attn_mask
|
||||
attn = attn + attn_mask
|
||||
attn = torch.softmax(attn, dim=-1)
|
||||
out = torch.einsum("hqk,khd->qhd", attn, value)
|
||||
out = out.to(device).to(dtype)
|
||||
return out
|
||||
|
||||
head_size = query.shape[-1]
|
||||
num_query_heads = query.shape[1]
|
||||
num_kv_heads = value_cache.shape[1]
|
||||
num_input_tokens = query.shape[0]
|
||||
|
||||
num_q_per_kv = num_query_heads // num_kv_heads
|
||||
slopes = (
|
||||
alibi_slopes.view(num_query_heads, 1, 1)
|
||||
if alibi_slopes is not None
|
||||
else alibi_slopes
|
||||
)
|
||||
|
||||
for i in range(num_input_tokens):
|
||||
q = query[i].unsqueeze(0)
|
||||
block_table = block_tables[i]
|
||||
context_len = int(context_lens[i])
|
||||
|
||||
keys = []
|
||||
values = []
|
||||
for j in range(context_len):
|
||||
block_number = int(block_table[j // block_size])
|
||||
block_offset = j % block_size
|
||||
|
||||
k = key_cache[block_number, :, block_offset, :]
|
||||
keys.append(k)
|
||||
|
||||
v = value_cache[block_number, :, block_offset, :]
|
||||
values.append(v)
|
||||
keys = torch.stack(keys, dim=0)
|
||||
values = torch.stack(values, dim=0)
|
||||
if num_q_per_kv > 1:
|
||||
keys = torch.repeat_interleave(keys, num_q_per_kv, dim=1)
|
||||
values = torch.repeat_interleave(values, num_q_per_kv, dim=1)
|
||||
if alibi_slopes is not None:
|
||||
offsets = get_alibi_mask(
|
||||
num_query_heads, context_len, output.device, output.dtype
|
||||
)
|
||||
mask = offsets * slopes
|
||||
mask = mask.to(output.dtype)
|
||||
if window_left != -1:
|
||||
index = torch.ones_like(mask, dtype=torch.int32, device=mask.device)
|
||||
index[:, :, (context_len - 1 - window_left) :] = 0
|
||||
index = index.bool()
|
||||
mask.masked_fill_(index, float("-inf"))
|
||||
else:
|
||||
if window_left != -1:
|
||||
mask = torch.zeros([1, 1, context_len], dtype=q.dtype, device=q.device)
|
||||
index = torch.ones_like(mask, dtype=torch.int32, device=mask.device)
|
||||
index[:, :, (context_len - 1 - window_left) :] = 0
|
||||
index = index.bool()
|
||||
mask.masked_fill_(index, float("-inf"))
|
||||
else:
|
||||
mask = None
|
||||
|
||||
out = ref_masked_attention(
|
||||
q,
|
||||
keys,
|
||||
values,
|
||||
scale,
|
||||
mask,
|
||||
)
|
||||
out = out.view(num_query_heads, head_size)
|
||||
if softcap != 0.0:
|
||||
out = softcap * torch.tanh(out / softcap)
|
||||
output[i].copy_(out, non_blocking=True)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def lmdeploy_paged_attention(
|
||||
output: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
num_kv_heads: torch.Tensor,
|
||||
scale: float,
|
||||
block_tables: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
block_size: int,
|
||||
max_context_len: int,
|
||||
alibi_slopes: torch.Tensor = None,
|
||||
softcap: float = 0.0,
|
||||
causal: bool = True,
|
||||
window_left: int = -1,
|
||||
window_right: int = -1,
|
||||
use_cuda_graph: bool = False,
|
||||
use_sqrt_alibi: bool = False,
|
||||
quant_type: int = 0,
|
||||
is_bbhh: bool = False,
|
||||
):
|
||||
"""
|
||||
is_bbhh = False key_cache, value_cache: [num_blocks, block_size, num_kv_heads, head_size]
|
||||
is_bbhh = True key_cache, value_cache: [num_blocks, num_kv_heads, block_size, head_size]
|
||||
|
||||
is_bbhh = False
|
||||
Arguments:
|
||||
query: [torch.half, torch.bfloat16] [num_tokens, num_heads, head_size]
|
||||
key_cache: [torch.half, torch.bfloat16] [num_blocks, num_kv_heads, block_size, head_size]
|
||||
value_cache: [torch.half, torch.bfloat16] [num_blocks, num_kv_heads, block_size, head_size]
|
||||
num_kv_heads: int
|
||||
scale: float
|
||||
block_tables: [torch.int64] [num_tokens, max_num_blocks_per_seq]
|
||||
context_lens: [torch.int32] [num_tokens]
|
||||
block_size: int
|
||||
max_context_len: int
|
||||
alibi_slopes: [torch.float32] [num_heads]
|
||||
softcap: float
|
||||
causal: bool
|
||||
window_left: int
|
||||
window_right: int
|
||||
use_sqrt_alibi: bool: False
|
||||
Return:
|
||||
output: [torch.half, torch.bfloat16] [num_tokens, num_heads, head_size]
|
||||
"""
|
||||
|
||||
ops.infer.lmdeploy_paged_attention(
|
||||
output,
|
||||
query,
|
||||
key_cache,
|
||||
value_cache,
|
||||
num_kv_heads,
|
||||
scale,
|
||||
block_tables,
|
||||
context_lens,
|
||||
block_size,
|
||||
max_context_len,
|
||||
alibi_slopes,
|
||||
causal,
|
||||
window_left,
|
||||
window_right,
|
||||
softcap,
|
||||
use_cuda_graph,
|
||||
use_sqrt_alibi,
|
||||
is_bbhh,
|
||||
quant_type,
|
||||
)
|
||||
return output
|
||||
|
||||
# lmdeploy_paged_attention = lmdeploy_paged_attention_ixinfer
|
||||
234
ixformer_sdk/inference/functions/marlin.py
Normal file
234
ixformer_sdk/inference/functions/marlin.py
Normal file
@@ -0,0 +1,234 @@
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"marlin_w4a16",
|
||||
"marlin_w4_weight_repack",
|
||||
"marlin_w8a16",
|
||||
"marlin_w8_weight_repack",
|
||||
]
|
||||
|
||||
|
||||
def marlin_w4a16(
|
||||
inputs: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
scales: torch.Tensor,
|
||||
zeros: torch.Tensor,
|
||||
bias: torch.Tensor = None, # TODO
|
||||
group_size: int = -1,
|
||||
format: str = "k16n32",
|
||||
batch_first: bool = True,
|
||||
outputs: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
inputs: (batch, m, k) if batch_first else (m, batch, k) torch.float16, torch.bfloat16
|
||||
weights: (batch, k/16, n/32, 64) torch.int32
|
||||
scales:
|
||||
(batch, k_groups, n) format:k16n32 torch.float16, torch.bfloat16
|
||||
(batch, n_groups, k) format:k16n32_grouped_n torch.float16, torch.bfloat16
|
||||
zeros:
|
||||
(batch, k_groups, n/8) format:k16n32 torch.int32
|
||||
(batch, n_groups, k/8) format:k16n32_grouped_n torch.int32
|
||||
group_size: int
|
||||
group size of quant
|
||||
format: str
|
||||
describe format of weight
|
||||
batch_first: bool
|
||||
describe format of input and output
|
||||
Returns:
|
||||
outputs: (batch, m, n) if batch_first else (m, batch, n) torch.float16, torch.bfloat16
|
||||
"""
|
||||
if outputs is None:
|
||||
batch, m = (
|
||||
(inputs.shape[0], inputs.shape[1])
|
||||
if batch_first
|
||||
else (inputs.shape[1], inputs.shape[0])
|
||||
)
|
||||
if format.startswith("k16n32"):
|
||||
n = weights.shape[2] * 32
|
||||
outputs = torch.empty(
|
||||
(batch, m, n) if batch_first else (m, batch, n),
|
||||
dtype=inputs.dtype,
|
||||
device=inputs.device,
|
||||
)
|
||||
|
||||
ops.infer.marlin_w4a16(
|
||||
outputs, inputs, weights, scales, zeros, bias, group_size, format, batch_first
|
||||
)
|
||||
return outputs
|
||||
|
||||
|
||||
def marlin_w4_weight_repack(
|
||||
weights: torch.Tensor,
|
||||
scales: torch.Tensor = None,
|
||||
zeros: torch.Tensor = None,
|
||||
weight_format: str = "gptq",
|
||||
reformat: str = "k16n32",
|
||||
pack_order: str = "default",
|
||||
repack_weight: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
weights:
|
||||
(batch, k, n/8) weight_format:awq torch.int32
|
||||
(batch, k/8, n) weight_format:gptq torch.int32
|
||||
scales:
|
||||
(batch, k_groups, n) format:k16n32 torch.float16, torch.bfloat16
|
||||
(batch, n_groups, k) format:k16n32_grouped_n torch.float16, torch.bfloat16
|
||||
zeros:
|
||||
(batch, k_groups, n/8) format:k16n32 torch.int32
|
||||
(batch, n_groups, k/8) format:k16n32_grouped_n torch.int32
|
||||
weight_format: str
|
||||
describe format of weight
|
||||
reformat: str
|
||||
describe format of repacked weight
|
||||
pack_order: str
|
||||
describe pack order on a pack unit
|
||||
Returns:
|
||||
repack_weight: (batch, k/16, n/32, 64) torch.int32
|
||||
"""
|
||||
assert weight_format in ["gptq", "gptq_grouped_n", "awq"]
|
||||
assert reformat in ["k16n32", "k16n32_grouped_n"]
|
||||
assert pack_order in ["default", "02461357", "01234567"]
|
||||
|
||||
if pack_order == "default":
|
||||
default_order = {
|
||||
"gptq": "01234567",
|
||||
"awq": "02461357",
|
||||
"gptq_grouped_n": "02461357",
|
||||
}
|
||||
pack_order = default_order[weight_format]
|
||||
|
||||
if weight_format.startswith("gptq"):
|
||||
batch, pack_k, n = weights.shape
|
||||
k = pack_k * 8
|
||||
elif weight_format == "awq":
|
||||
batch, k, pack_n = weights.shape
|
||||
n = pack_n * 8
|
||||
|
||||
repack_scales, repack_zeros = None, None
|
||||
if reformat.startswith("k16n32"):
|
||||
if repack_weight is None:
|
||||
repack_weight = torch.empty(
|
||||
(batch, k // 16, n // 32, 64),
|
||||
dtype=torch.int32,
|
||||
device=weights.device,
|
||||
)
|
||||
if scales is not None:
|
||||
repack_scales = torch.empty_like(scales)
|
||||
if zeros is not None:
|
||||
repack_zeros = torch.empty_like(zeros)
|
||||
|
||||
ops.infer.marlin_w4_weight_repack(
|
||||
weights,
|
||||
repack_weight,
|
||||
scales,
|
||||
repack_scales,
|
||||
zeros,
|
||||
repack_zeros,
|
||||
weight_format,
|
||||
reformat,
|
||||
pack_order,
|
||||
)
|
||||
|
||||
if repack_scales is not None and repack_zeros is not None:
|
||||
return repack_weight, repack_scales, repack_zeros
|
||||
else:
|
||||
return repack_weight
|
||||
|
||||
|
||||
def marlin_w8a16(
|
||||
inputs: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
scales: torch.Tensor,
|
||||
bias: torch.Tensor = None, # TODO
|
||||
group_size: int = -1,
|
||||
format: str = "k16n16",
|
||||
batch_first: bool = True,
|
||||
outputs: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
inputs: (batch, m, k) if batch_first else (m, batch, k) torch.float16, torch.bfloat16
|
||||
weights: (batch, k/16, n/16, 64) torch.int32
|
||||
scales:
|
||||
(batch, k_groups, n) format:k16n16 torch.float32
|
||||
(batch, n_groups, k) format:k16n16_grouped_n torch.float32
|
||||
group_size: int
|
||||
group size of quant
|
||||
format: str
|
||||
describe format of weight
|
||||
batch_first: bool
|
||||
describe format of input and output
|
||||
Returns:
|
||||
outputs: (batch, m, n) if batch_first else (m, batch, n) torch.float16, torch.bfloat16
|
||||
"""
|
||||
if outputs is None:
|
||||
batch, m = (
|
||||
(inputs.shape[0], inputs.shape[1])
|
||||
if batch_first
|
||||
else (inputs.shape[1], inputs.shape[0])
|
||||
)
|
||||
if format.startswith("k16n16"):
|
||||
n = weights.shape[2] * 16
|
||||
outputs = torch.empty(
|
||||
(batch, m, n) if batch_first else (m, batch, n),
|
||||
dtype=inputs.dtype,
|
||||
device=inputs.device,
|
||||
)
|
||||
|
||||
ops.infer.marlin_w8a16(
|
||||
outputs, inputs, weights, scales, bias, group_size, format, batch_first
|
||||
)
|
||||
return outputs
|
||||
|
||||
|
||||
def marlin_w8_weight_repack(
|
||||
weights: torch.Tensor,
|
||||
scales: torch.Tensor = None,
|
||||
weight_format: str = "int8",
|
||||
reformat: str = "k16n16",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
weights:
|
||||
(batch, k, n) weight_format:int8 torch.int8
|
||||
scales:
|
||||
(batch, k_groups, n) format:k16n16 torch.float32
|
||||
(batch, n_groups, k) format:k16n16_grouped_n torch.float32
|
||||
weight_format: str
|
||||
describe format of weight
|
||||
reformat: str
|
||||
describe format of repacked weight
|
||||
Returns:
|
||||
repack_weight:
|
||||
(batch, k/16, n/16, 64) torch.int32
|
||||
"""
|
||||
assert weight_format in ["int8"]
|
||||
assert reformat in ["k16n16", "k16n16_grouped_n"]
|
||||
|
||||
repack_scales = None
|
||||
if weight_format == "int8":
|
||||
batch, k, n = weights.shape
|
||||
repack_weight = torch.empty(
|
||||
(batch, k // 16, n // 16, 64),
|
||||
dtype=torch.int32,
|
||||
device=weights.device,
|
||||
)
|
||||
if scales is not None:
|
||||
repack_scales = torch.empty_like(scales)
|
||||
|
||||
ops.infer.marlin_w8_weight_repack(
|
||||
weights,
|
||||
repack_weight,
|
||||
scales,
|
||||
repack_scales,
|
||||
weight_format,
|
||||
reformat,
|
||||
)
|
||||
|
||||
if repack_scales is not None:
|
||||
return repack_weight, repack_scales
|
||||
else:
|
||||
return repack_weight
|
||||
50
ixformer_sdk/inference/functions/matmul.py
Normal file
50
ixformer_sdk/inference/functions/matmul.py
Normal file
@@ -0,0 +1,50 @@
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = ["matmul", "ref_matmul"]
|
||||
|
||||
|
||||
def ref_matmul(input, other, *, transa, transb, alpha):
|
||||
if transa:
|
||||
dims = list(range(input.ndim))
|
||||
dims[-1], dims[-2] = dims[-2], dims[-1]
|
||||
input = input.permute(*dims).contiguous()
|
||||
|
||||
if transb:
|
||||
dims = list(range(other.ndim))
|
||||
dims[-1], dims[-2] = dims[-2], dims[-1]
|
||||
other = other.permute(*dims).contiguous()
|
||||
|
||||
return alpha * torch.matmul(input, other)
|
||||
|
||||
|
||||
def matmul(
|
||||
input: torch.Tensor,
|
||||
other: torch.Tensor,
|
||||
*,
|
||||
transa: bool = False,
|
||||
transb: bool = False,
|
||||
alpha: float = 1.0,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Args:
|
||||
input: (...,m,k) or (...,k,m) torch.half
|
||||
当transa为False shape : [...,m,k], 当transa为True shape : [...,k,m]
|
||||
other: (...,k,n) or (...,n,k) torch.half
|
||||
当transa为False shape : [...,m,k], 当transa为True shape : [...,k,m]
|
||||
transa: bool
|
||||
transb: bool
|
||||
alpha: float
|
||||
Returns:
|
||||
Tensor: (..., m, n) torch.half
|
||||
"""
|
||||
if not input.is_contiguous():
|
||||
input = input.contiguous()
|
||||
|
||||
if not other.is_contiguous():
|
||||
if not other.transpose(-2, -1).is_contiguous():
|
||||
other = other.contiguous()
|
||||
|
||||
return ops.train.matmul(
|
||||
input, other, transa=transa, transb=transb, alpha=alpha, beta=0.0
|
||||
)
|
||||
325
ixformer_sdk/inference/functions/mla_fused.py
Normal file
325
ixformer_sdk/inference/functions/mla_fused.py
Normal file
@@ -0,0 +1,325 @@
|
||||
from typing import Optional
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
# 0.6.3
|
||||
"ref_minicpm3_fused_rope",
|
||||
"ref_minicpm3_fused_copy_kv",
|
||||
"minicpm3_fused_rope",
|
||||
"minicpm3_fused_copy_kv",
|
||||
# 0.6.6
|
||||
"ref_mla_rope_phi",
|
||||
"mla_rope_phi",
|
||||
"ref_mla_rope",
|
||||
"mla_rope",
|
||||
"ref_mla_copy_kv",
|
||||
"mla_copy_kv",
|
||||
]
|
||||
|
||||
|
||||
def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
|
||||
x1 = x[..., : x.shape[-1] // 2]
|
||||
x2 = x[..., x.shape[-1] // 2 :]
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
|
||||
|
||||
def _rotate_gptj(x: torch.Tensor) -> torch.Tensor:
|
||||
x1 = x[..., ::2]
|
||||
x2 = x[..., 1::2]
|
||||
x = torch.stack((-x2, x1), dim=-1)
|
||||
return x.flatten(-2)
|
||||
|
||||
|
||||
# vllm 0.6.3
|
||||
def ref_minicpm3_fused_rope(
|
||||
positions: torch.Tensor,
|
||||
long_prompt_offset: torch.Tensor,
|
||||
long_short_cos_sin_cache: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
out_query: Optional[torch.Tensor] = None,
|
||||
out_key: Optional[torch.Tensor] = None,
|
||||
):
|
||||
idx = torch.add(positions, long_prompt_offset)
|
||||
cos_sin = torch.index_select(long_short_cos_sin_cache, 0, idx)
|
||||
|
||||
cos, sin = cos_sin.chunk(2, dim=-1)
|
||||
cos = cos.repeat(1, 2).unsqueeze(-2)
|
||||
sin = sin.repeat(1, 2).unsqueeze(-2)
|
||||
|
||||
out_query = query * cos + _rotate_neox(query) * sin
|
||||
out_key = key * cos + _rotate_neox(key) * sin
|
||||
|
||||
return out_query, out_key
|
||||
|
||||
|
||||
def minicpm3_fused_rope(
|
||||
positions: torch.Tensor,
|
||||
long_prompt_offset: torch.Tensor,
|
||||
long_short_cos_sin_cache: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
out_query: Optional[torch.Tensor] = None,
|
||||
out_key: Optional[torch.Tensor] = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
positions: (num_tokens,) torch.int64
|
||||
long_prompt_offset: (num_tokens,) torch.int64
|
||||
long_short_cos_sin_cache: (max_length, head_dim) torch.float16, torch.bfloat16
|
||||
query: (num_tokens, num_q_heads, head_dim) same as long_short_cos_sin_cache
|
||||
key: (num_tokens, num_kv_heads, head_dim) same as long_short_cos_sin_cache
|
||||
out_query: same as query
|
||||
out_key: same as key
|
||||
Returns:
|
||||
out_query: same as query
|
||||
out_key: same as key
|
||||
"""
|
||||
|
||||
if out_query is None:
|
||||
out_query = torch.empty_like(query)
|
||||
if out_key is None:
|
||||
out_key = torch.empty_like(key)
|
||||
|
||||
ops.infer.minicpm3_fused_rope(
|
||||
positions,
|
||||
long_prompt_offset,
|
||||
long_short_cos_sin_cache,
|
||||
query,
|
||||
key,
|
||||
out_query,
|
||||
out_key,
|
||||
)
|
||||
return out_query, out_key
|
||||
|
||||
|
||||
def ref_minicpm3_fused_copy_kv(
|
||||
k_nope: torch.Tensor,
|
||||
k_pe: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
new_k: Optional[torch.Tensor] = None,
|
||||
new_v: Optional[torch.Tensor] = None,
|
||||
):
|
||||
num_tokens, num_heads, k_head_dim = k_nope.shape
|
||||
head_dim = k_pe.shape[-1] + k_head_dim
|
||||
v_head_dim = v.shape[-1]
|
||||
|
||||
if new_k is None:
|
||||
new_k = k_nope.new_empty([num_tokens, num_heads, head_dim])
|
||||
if new_v is None:
|
||||
new_v = k_nope.new_empty([num_tokens, num_heads, head_dim])
|
||||
|
||||
new_k[:, :, :k_head_dim] = k_nope
|
||||
new_k[:, :, k_head_dim:] = k_pe
|
||||
new_v[:, :, :v_head_dim] = v
|
||||
new_v[:, :, v_head_dim:] = 0
|
||||
|
||||
return new_k.view(num_tokens, -1), new_v.view(num_tokens, -1)
|
||||
|
||||
|
||||
def minicpm3_fused_copy_kv(
|
||||
k_nope: torch.Tensor,
|
||||
k_pe: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
new_k: Optional[torch.Tensor] = None,
|
||||
new_v: Optional[torch.Tensor] = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
k_nope: (num_tokens, num_heads, k_head_dim) torch.float16, torch.bfloat16
|
||||
k_pe: (num_tokens, 1, head_dim - k_head_dim) same as k_nope
|
||||
v: (num_tokens, num_heads, v_head_dim) same as k_nope
|
||||
new_k: (num_tokens, num_heads, head_dim) same as k_nope
|
||||
new_v: (num_tokens, num_heads, head_dim) same as k_nope
|
||||
Returns:
|
||||
new_k: (num_tokens, num_heads, head_dim) same as k_nope
|
||||
new_v: (num_tokens, num_heads, head_dim) same as k_nope
|
||||
"""
|
||||
|
||||
num_tokens, num_heads, k_head_dim = k_nope.shape
|
||||
head_dim = k_pe.shape[-1] + k_head_dim
|
||||
|
||||
if new_k is None:
|
||||
new_k = k_nope.new_empty([num_tokens, num_heads * head_dim])
|
||||
if new_v is None:
|
||||
new_v = k_nope.new_empty([num_tokens, num_heads * head_dim])
|
||||
|
||||
ops.infer.minicpm3_fused_copy_kv(k_nope, k_pe, v, new_k, new_v)
|
||||
|
||||
return new_k, new_v
|
||||
|
||||
|
||||
# vllm 0.6.6
|
||||
def ref_mla_rope_phi(
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
long_short_cos_sin_cache: torch.Tensor,
|
||||
k: int,
|
||||
offsets: Optional[torch.Tensor] = None,
|
||||
):
|
||||
long_prompt_offset = (
|
||||
torch.any(positions > k).float() * torch.full_like(positions, k)
|
||||
).long()
|
||||
idx = (
|
||||
torch.add(positions, long_prompt_offset)
|
||||
if long_prompt_offset is not None
|
||||
else positions
|
||||
)
|
||||
|
||||
idx = torch.add(idx, offsets) if offsets is not None else idx
|
||||
cos_sin = torch.index_select(long_short_cos_sin_cache, 0, idx)
|
||||
|
||||
cos, sin = cos_sin.chunk(2, dim=-1)
|
||||
cos = cos.repeat(1, 2).unsqueeze(-2)
|
||||
sin = sin.repeat(1, 2).unsqueeze(-2)
|
||||
|
||||
query = query * cos + _rotate_neox(query) * sin
|
||||
key = key * cos + _rotate_neox(key) * sin
|
||||
|
||||
return query, key
|
||||
|
||||
|
||||
def mla_rope_phi(
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
key_out: torch.Tensor,
|
||||
long_short_cos_sin_cache: torch.Tensor,
|
||||
long_offset: torch.Tensor,
|
||||
k: int,
|
||||
offsets: Optional[torch.Tensor] = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
positions: (num_tokens,) torch.int64
|
||||
query: (num_tokens, num_q_heads, head_dim) same as long_short_cos_sin_cache
|
||||
key: (num_tokens, 1, head_dim) same as long_short_cos_sin_cache
|
||||
key_out: (num_tokens, num_q_heads, head_dim) same as long_short_cos_sin_cache
|
||||
long_short_cos_sin_cache: (max_length, head_dim) same as long_short_cos_sin_cache
|
||||
long_offset: (1,) torch.bool
|
||||
k: int
|
||||
offsets: (num_tokens,)
|
||||
Returns:
|
||||
query:
|
||||
key_out:
|
||||
"""
|
||||
|
||||
ops.infer.mla_rope_phi(
|
||||
positions,
|
||||
query,
|
||||
key,
|
||||
key_out,
|
||||
long_short_cos_sin_cache,
|
||||
long_offset,
|
||||
k,
|
||||
offsets,
|
||||
)
|
||||
return query, key_out
|
||||
|
||||
|
||||
def ref_mla_rope(
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
offsets: Optional[torch.Tensor] = None,
|
||||
rotary_dim: int = None,
|
||||
is_neox_style: bool = False,
|
||||
):
|
||||
"""PyTorch-native implementation equivalent to forward()."""
|
||||
head_size = query.size(-1)
|
||||
rotary_dim = rotary_dim or head_size
|
||||
query_rot = query[..., :rotary_dim]
|
||||
key_rot = key[..., :rotary_dim]
|
||||
if rotary_dim < head_size:
|
||||
query_pass = query[..., rotary_dim:]
|
||||
key_pass = key[..., rotary_dim:]
|
||||
|
||||
cos_sin = cos_sin_cache[
|
||||
torch.add(positions, offsets) if offsets is not None else positions
|
||||
]
|
||||
cos, sin = cos_sin.chunk(2, dim=-1)
|
||||
if is_neox_style:
|
||||
cos = cos.repeat(1, 1, 2).unsqueeze(-2)
|
||||
sin = sin.repeat(1, 1, 2).unsqueeze(-2)
|
||||
else:
|
||||
cos = cos.repeat_interleave(2, dim=-1).unsqueeze(-2)
|
||||
sin = sin.repeat_interleave(2, dim=-1).unsqueeze(-2)
|
||||
|
||||
rotate_fn = _rotate_neox if is_neox_style else _rotate_gptj
|
||||
query_rot = query_rot * cos + rotate_fn(query_rot) * sin
|
||||
key_rot = key_rot * cos + rotate_fn(key_rot) * sin
|
||||
|
||||
if rotary_dim < head_size:
|
||||
query = torch.cat((query_rot, query_pass), dim=-1)
|
||||
key = torch.cat((key_rot, key_pass), dim=-1)
|
||||
else:
|
||||
query = query_rot
|
||||
key = key_rot
|
||||
return query, key
|
||||
|
||||
|
||||
def mla_rope(
|
||||
positions: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
key_out: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
offsets: Optional[torch.Tensor] = None,
|
||||
is_neox_style: bool = False,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
positions: (num_tokens,) torch.int64
|
||||
query: (num_tokens, num_q_heads, head_dim) torch.half torch.bfloat torch.float
|
||||
key: (num_tokens, 1, head_dim) same as query
|
||||
key_out: (num_tokens, num_q_heads, head_dim) same as query
|
||||
cos_sin_cache: (max_length, head_dim) same as query
|
||||
offsets: (num_tokens,) same as query
|
||||
is_neox_style: bool
|
||||
Returns:
|
||||
query:
|
||||
key_out:
|
||||
"""
|
||||
|
||||
ops.infer.mla_rope(
|
||||
positions,
|
||||
query,
|
||||
key,
|
||||
key_out,
|
||||
cos_sin_cache,
|
||||
is_neox_style,
|
||||
offsets,
|
||||
)
|
||||
return query, key_out
|
||||
|
||||
|
||||
def ref_mla_copy_kv(key_pe, key_nope, value_nope):
|
||||
shape = key_nope.shape[:-1] + (key_pe.shape[-1] + key_nope.shape[-1],)
|
||||
key = torch.empty(shape, device=key_nope.device, dtype=key_nope.dtype)
|
||||
value = torch.empty_like(key)
|
||||
|
||||
key[..., : key_nope.size(-1)] = key_nope
|
||||
key[..., key_nope.size(-1) :] = key_pe
|
||||
value[..., : value_nope.size(-1)] = value_nope
|
||||
value[..., value_nope.size(-1) :] = 0.0
|
||||
return key, value
|
||||
|
||||
|
||||
def mla_copy_kv(key_nope, value_nope, key, value):
|
||||
"""
|
||||
Args:
|
||||
key_nope: (num_tokens, num_heads, k_nope_dim) torch.float16, torch.bfloat16, torch.float
|
||||
value_nope: (num_tokens, num_heads, v_head_dim) same as key_nope
|
||||
key: (num_tokens, num_heads, head_dim) same as key_nope
|
||||
value: (num_tokens, num_heads, head_dim) same as key_nope
|
||||
Returns:
|
||||
key:
|
||||
value:
|
||||
"""
|
||||
|
||||
ops.infer.mla_copy_kv(key_nope, value_nope, key, value)
|
||||
return key, value
|
||||
315
ixformer_sdk/inference/functions/mm.py
Normal file
315
ixformer_sdk/inference/functions/mm.py
Normal file
@@ -0,0 +1,315 @@
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
import torch.nn.functional
|
||||
|
||||
__all__ = [
|
||||
"mm",
|
||||
"addmm",
|
||||
"fused_addmm_bias_col_act",
|
||||
"ref_fused_addmm_bias_col_act",
|
||||
"ref_addmm",
|
||||
"ref_mm",
|
||||
"ref_bmm",
|
||||
"bmm",
|
||||
]
|
||||
|
||||
|
||||
def ref_mm(input, mat, *, out=None):
|
||||
out = torch.mm(input, mat, out = out)
|
||||
return out
|
||||
|
||||
|
||||
def mm(input, mat, *, out=None):
|
||||
|
||||
"""
|
||||
Args:
|
||||
input: (m,k) torch.float16, torch.bfloat16, torch.float32
|
||||
mat: (k,n) torch.float16, torch.bfloat16, torch.float32
|
||||
out: (m,n) torch.float16, torch.bfloat16, torch.float32
|
||||
Returns:
|
||||
out: (m,n) torch.float16, torch.bfloat16, torch.float32
|
||||
"""
|
||||
assert input.dim() == mat.dim(), "mm tensors must be 2-D"
|
||||
assert input.size(1) == mat.size(
|
||||
0
|
||||
), f"mm cannot be multiplied, {input.size(0)}X{input.size(1)} and {mat.size(0)}X{mat.size(1)}"
|
||||
|
||||
m = input.shape[0]
|
||||
n = mat.shape[-1]
|
||||
if out is None:
|
||||
out = input.new_empty([m, n])
|
||||
ops.infer.mm(input, mat, out)
|
||||
return out
|
||||
|
||||
|
||||
"""ixinfer support activations
|
||||
/// @ingroup GEMM
|
||||
typedef enum {
|
||||
CUINFER_BLAS_GEMM_CUSTOM_NONE = 0,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_BIAS_ADD_ROW_OUT = 1,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS = 2,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_GELU = 3,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_RELU = 4,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_TRANSPOSE = 5,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS = 6,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS_GELU = 7,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS_RELU = 8,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS_TRANSPOSE = 9,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_SIGMOID = 10,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS_SIGMOID = 11,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_SILU = 12,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS_SILU = 13,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_SIGMOID = 14,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_SILU = 15,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_TANH = 16,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_FLOATBIAS_TANH = 17,
|
||||
CUINFER_BLAS_GEMM_SPECIAL_INT8_FLOATBIAS = 18,
|
||||
CUINFER_BLAS_GEMM_SPECIAL_INT8_FLOATBIAS_GELU = 19,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_SWISH = 20,
|
||||
CUINFER_BLAS_GEMM_CUSTOM_HALFBIAS_ERF_GELU = 21
|
||||
} cuinferGEMMCustomOption_t;
|
||||
"""
|
||||
|
||||
activation_to_id = {
|
||||
"fused_bias_col": 2, # support bf16, fp16
|
||||
"fused_bias_gelu": 3, # support fp16
|
||||
"fused_bias_relu": 4, # support fp16
|
||||
}
|
||||
|
||||
id_to_activation = {value: key for key, value in activation_to_id.items()}
|
||||
|
||||
|
||||
def ref_addmm(input, mat1, mat2, *, beta=1, alpha=1, out=None):
|
||||
output_pt = torch.addmm(input, mat1, mat2, alpha=alpha, beta=beta)
|
||||
return output_pt
|
||||
|
||||
|
||||
def ref_fused_addmm_bias_col_act(
|
||||
input, mat1, mat2, *, beta=1, alpha=1, out=None, bias=None, activation=2
|
||||
):
|
||||
if isinstance(activation, int):
|
||||
assert activation in id_to_activation
|
||||
if isinstance(activation, str):
|
||||
assert activation in activation_to_id
|
||||
activation = activation_to_id[activation]
|
||||
if activation == 2:
|
||||
output_pt = torch.addmm(input, mat1, mat2, alpha=alpha, beta=beta) + bias
|
||||
elif activation == 3:
|
||||
output_pt = torch.nn.functional.gelu(
|
||||
torch.addmm(input, mat1, mat2, alpha=alpha, beta=beta) + bias
|
||||
)
|
||||
else:
|
||||
output_pt = torch.nn.functional.relu(
|
||||
torch.addmm(input, mat1, mat2, alpha=alpha, beta=beta) + bias
|
||||
)
|
||||
return output_pt
|
||||
|
||||
|
||||
def addmm(input, mat1, mat2, *, beta=1, alpha=1, out=None):
|
||||
|
||||
"""
|
||||
Args:
|
||||
input: (m,n) torch.float16, torch.bfloat16, torch.float32
|
||||
mat1: (m,k) torch.float16, torch.bfloat16, torch.float32
|
||||
mat2: (k,n) torch.float16, torch.bfloat16, torch.float32
|
||||
beta: float
|
||||
alpha: float
|
||||
out: (m,n) torch.float16, torch.bfloat16, torch.float32
|
||||
Returns:
|
||||
out: (m,n) torch.float16, torch.bfloat16, torch.float32
|
||||
"""
|
||||
assert mat1.dim() == mat2.dim(), "addmm mat1 mat2 tensors must be 2-D"
|
||||
assert mat1.size(1) == mat2.size(
|
||||
0
|
||||
), f"addmm cannot be multiplied, {mat1.size(0)}X{mat1.size(1)} and {mat2.size(0)}X{mat2.size(1)}"
|
||||
|
||||
m = mat1.shape[0]
|
||||
n = mat2.shape[-1]
|
||||
|
||||
if input is not None and len(input.shape) == 1:
|
||||
input = input.view(1, -1)
|
||||
|
||||
if out is None:
|
||||
out = input.new_empty([m, n])
|
||||
if input is None:
|
||||
input = out
|
||||
beta = 0
|
||||
|
||||
ops.infer.addmm(input, mat1, mat2, beta, alpha, out, None, 0)
|
||||
return out
|
||||
|
||||
|
||||
def fused_addmm_bias_col_act(
|
||||
input, mat1, mat2, *, beta=1, alpha=1, out=None, bias=None, activation=2
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (m,n) torch.float16, torch.bfloat16, torch.float32
|
||||
mat1: (m,k) torch.float16, torch.bfloat16, torch.float32
|
||||
mat2: (k,n) torch.float16, torch.bfloat16, torch.float32
|
||||
beta: float
|
||||
alpha: float
|
||||
out: (m,n) torch.float16, torch.bfloat16, torch.float32
|
||||
当out的shape为(m,n)时,如果out是is_continouns,则bias 必须为(1,n), 否则,bias为(m,1)
|
||||
bias: (1,n) or (m,1)
|
||||
activation: str or int
|
||||
"fused_bias_col": 2, "fused_bias_gelu": 3, "fused_bias_relu": 4
|
||||
Returns:
|
||||
out: (m,n) torch.float16, torch.bfloat16, torch.float32
|
||||
"""
|
||||
assert mat1.dim() == mat2.dim(), "addmm mat1 mat2 tensors must be 2-D"
|
||||
assert mat1.size(1) == mat2.size(
|
||||
0
|
||||
), f"addmm cannot be multiplied, {mat1.size(0)}X{mat1.size(1)} and {mat2.size(0)}X{mat2.size(1)}"
|
||||
|
||||
if isinstance(activation, int):
|
||||
assert activation in id_to_activation
|
||||
if isinstance(activation, str):
|
||||
assert activation in activation_to_id
|
||||
activation = activation_to_id[activation]
|
||||
|
||||
m = mat1.shape[0]
|
||||
n = mat2.shape[-1]
|
||||
|
||||
if out is None:
|
||||
out = input.new_empty([m, n])
|
||||
if input is None:
|
||||
input = out
|
||||
beta = 0
|
||||
|
||||
# activations
|
||||
if activation == 2:
|
||||
assert bias is not None
|
||||
if mat1.dtype == torch.float:
|
||||
ops.infer.addmm(input, mat1, mat2, beta, alpha, out, None, 0)
|
||||
out.add_(bias)
|
||||
return out
|
||||
elif activation == 3:
|
||||
assert bias is not None
|
||||
if mat1.dtype == torch.float:
|
||||
ops.infer.addmm(input, mat1, mat2, beta, alpha, out, None, 0)
|
||||
out.add_(bias)
|
||||
out.copy_(torch.nn.functional.gelu(out))
|
||||
return out
|
||||
elif mat1.dtype == torch.bfloat16:
|
||||
ops.infer.addmm(input, mat1, mat2, beta, alpha, out, bias, 2)
|
||||
out.copy_(torch.nn.functional.gelu(out))
|
||||
return out
|
||||
elif activation == 4:
|
||||
assert bias is not None
|
||||
if mat1.dtype == torch.float:
|
||||
ops.infer.addmm(input, mat1, mat2, beta, alpha, out, None, 0)
|
||||
out.add_(bias)
|
||||
out.copy_(torch.nn.functional.relu(out))
|
||||
return out
|
||||
elif mat1.dtype == torch.bfloat16:
|
||||
ops.infer.addmm(input, mat1, mat2, beta, alpha, out, bias, 2)
|
||||
out.copy_(torch.nn.functional.relu(out))
|
||||
return out
|
||||
|
||||
ops.infer.addmm(input, mat1, mat2, beta, alpha, out, bias, activation)
|
||||
return out
|
||||
|
||||
|
||||
def ref_bmm(
|
||||
input: torch.Tensor,
|
||||
mat2: torch.Tensor,
|
||||
alpha: float = 1,
|
||||
format: str = "NN",
|
||||
input_scales: torch.Tensor = None,
|
||||
mat2_scales: torch.Tensor = None,
|
||||
out_dtype: torch.dtype = None,
|
||||
out: torch.Tensor = None,
|
||||
):
|
||||
if format[1] == "T":
|
||||
input = input.transpose(-1, -2)
|
||||
if format[0] == "T":
|
||||
mat2 = mat2.transpose(-1, -2)
|
||||
|
||||
m = input.size(-2)
|
||||
n = mat2.size(-1)
|
||||
|
||||
bs = input.size(0)
|
||||
if input.dtype != torch.int8:
|
||||
out_dtype = input.dtype
|
||||
if out is None:
|
||||
out = torch.empty([bs, m, n], dtype=out_dtype, device=input.device)
|
||||
|
||||
if input.dtype != torch.int8:
|
||||
torch.bmm(input, mat2, out=out)
|
||||
if alpha != 1:
|
||||
out = out * alpha
|
||||
else:
|
||||
input = input.float() * input_scales.view(1, -1, 1)
|
||||
mat2 = mat2.float() * mat2_scales.view(1, 1, -1)
|
||||
out = torch.bmm(input.float(), mat2.float()) * alpha
|
||||
out = out.to(out_dtype)
|
||||
return out
|
||||
|
||||
|
||||
def bmm(
|
||||
input: torch.Tensor,
|
||||
mat2: torch.Tensor,
|
||||
alpha: float = 1,
|
||||
format: str = "NN",
|
||||
input_scales: torch.Tensor = None,
|
||||
mat2_scales: torch.Tensor = None,
|
||||
out_dtype: torch.dtype = None,
|
||||
out: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
out = (input@mat2)*alpha
|
||||
Support three formats:
|
||||
format: "NN" input shape: (b, m, k) mat2 shape: (b, k, n) out shape: (b, m, n).
|
||||
If the dtype of input is int8, the following conditions need to be met: n%64==0 k%64==0
|
||||
format: "TN" input shape: (b, m, k) mat2 shape: (b, n, k) out shape: (b, m, n)
|
||||
If the dtype of input is int8, the following conditions need to be met: n%2==0 k%64==0
|
||||
format: "NT" input shape: (b, k, m) mat2 shape: (b, k, n) out shape: (b, m, n)
|
||||
If the dtype of input is int8, the following conditions need to be met: m%64==0 n%64==0 k%64==0
|
||||
If the dtype of input is int8, it is necessary to specify out_dtype.
|
||||
Args:
|
||||
input: (b, m, k) or (b, k, m) torch.float16, torch.bfloat16, int8
|
||||
mat2: (b, k, n) or (b, n, k) torch.float16, torch.bfloat16, int8
|
||||
alpha: float32
|
||||
format: TN,NN,NT string
|
||||
input_scales: (m) torch.float32
|
||||
mat2_scales: (n) torch.float32
|
||||
out_dtype: torch.float16, torch.bfloat16
|
||||
out: (b, m, n) torch.float16, torch.bfloat16
|
||||
Returns:
|
||||
out: (b, m, n) torch.float16, torch.bfloat16
|
||||
"""
|
||||
|
||||
if format[1] == "N":
|
||||
m = input.size(-2)
|
||||
k = input.size(-1)
|
||||
else:
|
||||
m = input.size(-1)
|
||||
k = input.size(-2)
|
||||
if format[0] == "N":
|
||||
n = mat2.size(-1)
|
||||
else:
|
||||
n = mat2.size(-2)
|
||||
|
||||
if input.dtype == torch.int8:
|
||||
if format == "TN":
|
||||
assert (
|
||||
n % 2 == 0 and k % 64 == 0
|
||||
), f"bmm shape error, m={m} n={n} k={k}."
|
||||
elif format == "NT":
|
||||
assert (
|
||||
m % 64 == 0 and n % 64 == 0 and k % 64 == 0
|
||||
), f"bmm shape error, m={m} n={n} k={k}."
|
||||
elif format == "NN":
|
||||
assert (
|
||||
n % 64 == 0 and k % 64 == 0
|
||||
), f"bmm shape error, m={m} n={n} k={k}."
|
||||
bs = input.size(0)
|
||||
if out is None:
|
||||
if input.dtype != torch.int8:
|
||||
out_dtype = input.dtype
|
||||
else:
|
||||
assert out_dtype is not None
|
||||
out = torch.empty([bs, m, n], dtype=out_dtype, device=input.device)
|
||||
ops.infer.bmm(input, mat2, input_scales, mat2_scales, alpha, format, out)
|
||||
return out
|
||||
1380
ixformer_sdk/inference/functions/moe.py
Normal file
1380
ixformer_sdk/inference/functions/moe.py
Normal file
File diff suppressed because it is too large
Load Diff
84
ixformer_sdk/inference/functions/overlap_comm.py
Normal file
84
ixformer_sdk/inference/functions/overlap_comm.py
Normal file
@@ -0,0 +1,84 @@
|
||||
import itertools
|
||||
from functools import partial
|
||||
from typing import Callable, Dict, Iterable, Tuple
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
import ixformer.distributed as ixfd
|
||||
from ixformer.core.dispatcher import Dispatcher
|
||||
from ixformer.core.operator_autotuning import (
|
||||
OperatorPreBaseRangeAutotuning,
|
||||
sync_ranks_metric,
|
||||
)
|
||||
from ixformer.distributed import overlap_comm
|
||||
from ixformer.inference.overlap.linear_mlp_overlap_comm import linear_mlp_overlap
|
||||
from ixformer.distributed.overlap_comm import GemmMethod
|
||||
|
||||
__all__ = ["linear_allreduce_overlap", "linear_mlp_overlap"]
|
||||
|
||||
|
||||
class LinearAllReducePreAutotuning(OperatorPreBaseRangeAutotuning, Dispatcher):
|
||||
def __init__(self, comm_group, *args, **kwargs):
|
||||
dist_barrier = True
|
||||
if "dist_barrier" in kwargs:
|
||||
dist_barrier = kwargs.pop("dist_barrier")
|
||||
|
||||
super().__init__(dist_barrier=dist_barrier, *args, **kwargs)
|
||||
self._comm_group = comm_group
|
||||
self._world_size = ixfd.get_group_world_size(comm_group)
|
||||
|
||||
@classmethod
|
||||
def dispatcher_key(cls, comm_group, *args, **kwargs):
|
||||
return (comm_group,)
|
||||
|
||||
def operators(self):
|
||||
chunks = [2, 4]
|
||||
gemm_algos = [GemmMethod.kCUINFER, GemmMethod.kCUBLAS, GemmMethod.kLIMITED_GEMM]
|
||||
candidate_ops = [overlap_comm.GemmAllReduceSplitOverlapComm.native_forward]
|
||||
for num_chunks, algo in itertools.product(chunks, gemm_algos):
|
||||
candidate_ops.append(
|
||||
partial(
|
||||
overlap_comm.linear_allreduce_overlap,
|
||||
num_chunks=num_chunks,
|
||||
gemm_method=algo,
|
||||
)
|
||||
)
|
||||
|
||||
return candidate_ops
|
||||
|
||||
@property
|
||||
def _gemm_shapes(self):
|
||||
basic_k = [4096, 6114, 8192]
|
||||
tp_k = [k // self._world_size for k in basic_k]
|
||||
basic_k = tp_k
|
||||
|
||||
basic_m = (512, 1024, 2048, 4096, 8192)
|
||||
|
||||
shapes = set(itertools.product(basic_m, basic_k))
|
||||
|
||||
return shapes
|
||||
|
||||
def get_operator_key(self, input, *args, **kwargs):
|
||||
ndim = input.ndim
|
||||
shape = input.shape
|
||||
|
||||
if ndim == 1:
|
||||
return (1, shape[0])
|
||||
elif ndim == 2:
|
||||
return shape
|
||||
else:
|
||||
return (sum(shape[:-1]), shape[-1])
|
||||
|
||||
def generate_operator_inputs(self) -> Iterable[Tuple[Tuple, Dict]]:
|
||||
for m, kn in self._gemm_shapes:
|
||||
input = torch.randn(m, kn, device="cuda", dtype=torch.half)
|
||||
weight = torch.randn(kn, kn, device="cuda", dtype=torch.half)
|
||||
yield (input, weight), {}
|
||||
|
||||
def perf_operator_time(self, op: Callable, *args, **kwargs) -> float:
|
||||
op_time = super().perf_operator_time(op, *args, **kwargs)
|
||||
return sync_ranks_metric(op_time, group=self._comm_group)
|
||||
|
||||
|
||||
linear_allreduce_overlap = overlap_comm.linear_allreduce_overlap
|
||||
143
ixformer_sdk/inference/functions/paged_attention.py
Normal file
143
ixformer_sdk/inference/functions/paged_attention.py
Normal file
@@ -0,0 +1,143 @@
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"paged_attention",
|
||||
"paged_attention_flashinfer",
|
||||
"paged_attention_cache_appended",
|
||||
]
|
||||
# paged_attention_cache_append
|
||||
|
||||
|
||||
def paged_attention_cache_appended(
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
slot_mapping: torch.Tensor,
|
||||
kv_cache_format: str = "HND", # STD NHD HND
|
||||
key_cache_scales: torch.Tensor = None,
|
||||
value_cache_scales: torch.Tensor = None,
|
||||
):
|
||||
if isinstance(key, torch.Tensor):
|
||||
ops.infer.paged_attention_cache_appended(
|
||||
key,
|
||||
value,
|
||||
key_cache,
|
||||
value_cache,
|
||||
slot_mapping,
|
||||
key.stride(0),
|
||||
value.stride(0),
|
||||
key_cache.stride(0),
|
||||
value_cache.stride(0),
|
||||
kv_cache_format,
|
||||
key_cache_scales,
|
||||
value_cache_scales,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def paged_attention(
|
||||
output: 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: torch.Tensor = None,
|
||||
use_sqrt_alibi: bool = False,
|
||||
key_cache_scales: torch.Tensor = None,
|
||||
value_cache_scales: torch.Tensor = None,
|
||||
kv_cache_format: str = "HND",
|
||||
algo: int = -1,
|
||||
):
|
||||
"""
|
||||
kv_cache_format
|
||||
STD : k/v format as same as vllm
|
||||
NHD : k/v format is [block_size, num_kv_heads, head_dim] in one page
|
||||
HND : k/v format is [num_kv_heads, block_size, head_dim] in one page
|
||||
algo
|
||||
-1 : auto chooes algorithm according to kv_cache_format
|
||||
0 : use the first algorithm
|
||||
1 : use the second algorithm
|
||||
"""
|
||||
if isinstance(query, torch.Tensor):
|
||||
ops.infer.paged_attention(
|
||||
output,
|
||||
query,
|
||||
key_cache,
|
||||
value_cache,
|
||||
num_kv_heads,
|
||||
scale,
|
||||
block_tables,
|
||||
seq_lens,
|
||||
block_size,
|
||||
max_seq_len,
|
||||
use_sqrt_alibi,
|
||||
alibi_slopes,
|
||||
key_cache_scales,
|
||||
value_cache_scales,
|
||||
kv_cache_format,
|
||||
algo,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def paged_attention_flashinfer(
|
||||
output: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
paged_kv_data,
|
||||
paged_kv_indptr: torch.Tensor,
|
||||
paged_kv_indices: torch.Tensor,
|
||||
paged_kv_last_page_len: torch.Tensor,
|
||||
scale: float,
|
||||
max_seq_len: int = -1,
|
||||
use_sqrt_alibi: bool = False,
|
||||
alibi_slopes: torch.Tensor = None,
|
||||
kv_cache_format: str = "HND",
|
||||
# key_cache_scales: torch.Tensor = None,
|
||||
# value_cache_scales: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
out / query : [num_seqs, num_qo_heads, head_size]
|
||||
paged_kv_data
|
||||
Tensor:
|
||||
NHD [max_num_pages, 2, page_size, num_kv_heads, head_size]
|
||||
HND [max_num_pages, 2, num_kv_heads, page_size, head_size]
|
||||
tuple(k_data, v_data)
|
||||
NHD [max_num_pages, page_size, num_kv_heads, head_size]
|
||||
HND [max_num_pages, num_kv_heads, page_size, head_size]
|
||||
paged_kv_indptr int32 : [num_seqs + 1]
|
||||
paged_kv_indices int32 : [max_num_pages]
|
||||
paged_kv_last_page_len int32 : [num_seqs]
|
||||
"""
|
||||
if isinstance(paged_kv_data, tuple):
|
||||
k_data, v_data = paged_kv_data
|
||||
pack_kv_data = (None, k_data, v_data)
|
||||
else:
|
||||
pack_kv_data = (paged_kv_data, None, None)
|
||||
|
||||
if isinstance(query, torch.Tensor):
|
||||
ops.infer.paged_attention_flashinfer(
|
||||
output,
|
||||
query,
|
||||
*pack_kv_data,
|
||||
paged_kv_indptr,
|
||||
paged_kv_indices,
|
||||
paged_kv_last_page_len,
|
||||
scale,
|
||||
max_seq_len,
|
||||
use_sqrt_alibi,
|
||||
alibi_slopes,
|
||||
kv_cache_format,
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
238
ixformer_sdk/inference/functions/quantized_linear.py
Normal file
238
ixformer_sdk/inference/functions/quantized_linear.py
Normal file
@@ -0,0 +1,238 @@
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
__all__ = [
|
||||
"quantized_linear",
|
||||
"quantized_weight_dequant",
|
||||
"ref_quantized_weight_dequant",
|
||||
"weight_quantize",
|
||||
]
|
||||
|
||||
|
||||
def quantized_linear(
|
||||
inputs: torch.Tensor,
|
||||
qweights: torch.Tensor,
|
||||
scales: torch.Tensor,
|
||||
quant_type: str,
|
||||
bits: int,
|
||||
qzeros: torch.Tensor = None,
|
||||
bias: torch.Tensor = None,
|
||||
group_size: int = -1,
|
||||
g_idx: torch.Tensor = None,
|
||||
format: str = "unknown",
|
||||
):
|
||||
"""
|
||||
QuantType inputs qweights Scales bits qzeros bias GroupSize Format ApiCall 备注
|
||||
awq (bs, ic) bf16/fp16 int32 NN:(ic, oc // 8) TN:(oc, ic // 8) (ic // group_size, oc)fp16/bf16 4/8 int32(ic // group_size, oc // 8) (oc) or None fp16/bf16 32/128 TN/NN vllm & auto-awq
|
||||
gptq (bs, ic) bf16/fp16 int32 (ic//8, oc) (ic // group_size, oc)fp16/bf16 4 int32(ic // group_size, oc // 8) (oc) or None fp16/bf16 ic/128 \ auto-gptq bs 只支持到8
|
||||
fp4 (bs, ic) bf16/fp16 uint8 (oc * ic // 2, 1) (oc * ic // group_size)fp32 4 \ (oc) or None fp16/bf16 64 \ bitsandbytes bs 只支持到8
|
||||
nf4 (bs, ic) bf16/fp16 uint8 (oc * ic // 2, 1) (oc * ic // group_size)fp32 4 \ (oc) or None fp16/bf16 64 \ bitsandbytes bs 只支持到8
|
||||
int8 (bs, ic) bf16/fp16 int8 TN:(oc, ic) NN:(ic, oc) (1, oc)fp16/bf16 8 \ (oc) or None fp16/bf16 -1 TN/NN vllm & bitsandbytes
|
||||
|
||||
"""
|
||||
if isinstance(inputs, torch.Tensor) and not inputs.requires_grad:
|
||||
return ops.infer.quantized_linear(
|
||||
inputs,
|
||||
qweights,
|
||||
scales,
|
||||
quant_type,
|
||||
bits,
|
||||
qzeros,
|
||||
bias,
|
||||
group_size,
|
||||
g_idx,
|
||||
format,
|
||||
)
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def quantized_weight_dequant(
|
||||
qweights: torch.Tensor,
|
||||
scales: torch.Tensor,
|
||||
quant_type: str,
|
||||
output_type: str,
|
||||
bits: int,
|
||||
qzeros: torch.Tensor = None,
|
||||
group_size: int = -1,
|
||||
g_idx: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
qweights: (oc, ic//2) or (ic// (32/bits, oc) torch.unint8 or torch.int32
|
||||
scales: (oc * ic//g) or (ic // g, oc) torch.float16, torch.bfloat16, torch.float32
|
||||
quant_type: str
|
||||
可选项:fp4/nf4/gptq/gptq-ex
|
||||
output_type: str
|
||||
可选项:fp16/bf16
|
||||
bits: int
|
||||
可选项:4/8
|
||||
qzeros: (ic//g, oc//(32/bits)) torch.int32
|
||||
group_size: int
|
||||
可选项:-1/64/128
|
||||
g_idx: (ic) torch.int
|
||||
|
||||
Returns:
|
||||
Tensor: (oc, ic) or (ic, oc) torch.float16, torch.bfloat16
|
||||
|
||||
quant_type qweights scales qzeros output_type bits group_size g_idx
|
||||
fp4/nf4 (oc, ic//2) uint8 (oc * ic//g) fp32 fp16/bf16 / 64/128 /
|
||||
gptq/gptq-ex (ic// (32/bits, oc) int32 (ic // g, oc) fp16/bf16 (ic // g, oc // (32/bits)) int32 fp16/bf16 4/8 -1 (ic)
|
||||
|
||||
"""
|
||||
if isinstance(qweights, torch.Tensor) and not qweights.requires_grad:
|
||||
return ops.infer.quantized_weight_dequant(
|
||||
qweights, scales, quant_type, output_type, bits, qzeros, group_size, g_idx
|
||||
)
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
def ref_quantized_weight_dequant(
|
||||
qweights: torch.Tensor,
|
||||
scales: torch.Tensor,
|
||||
quant_type: str,
|
||||
output_type: torch.dtype,
|
||||
bits: int,
|
||||
qzeros: torch.Tensor = None,
|
||||
group_size: int = -1,
|
||||
g_idx: torch.Tensor = None,
|
||||
order_map: list = None,
|
||||
):
|
||||
assert quant_type in ["awq"]
|
||||
if quant_type == "awq":
|
||||
# qweights:(k, n/8) int32
|
||||
# scale:(k/group_size, n) f16
|
||||
# qzeros:(k/group_size, n/8) int32
|
||||
ic, oc = qweights.shape[0], scales.shape[1]
|
||||
assert bits == 4
|
||||
if order_map is None:
|
||||
order_map = [0, 2, 4, 6, 1, 3, 5, 7]
|
||||
order_map = torch.Tensor(order_map).to(torch.int32).to(qweights.device)
|
||||
order_map = order_map.argsort()
|
||||
|
||||
# (1, 8)
|
||||
wf = (
|
||||
torch.tensor(list(range(0, 32, bits)), dtype=torch.int32)
|
||||
.unsqueeze(0)
|
||||
.to(qweights.device)
|
||||
)
|
||||
|
||||
# unpack qzeros
|
||||
unpack_zeros = torch.bitwise_right_shift(
|
||||
torch.unsqueeze(qzeros, 2).expand(-1, -1, 32 // bits), wf.unsqueeze(0)
|
||||
).to(torch.int16 if bits == 8 else torch.int8)
|
||||
unpack_zeros = unpack_zeros[:, :, order_map]
|
||||
unpack_zeros = torch.bitwise_and(unpack_zeros, (2**bits) - 1)
|
||||
# groups, 1, n
|
||||
unpack_zeros = unpack_zeros.reshape(unpack_zeros.shape[0], 1, -1)
|
||||
|
||||
# unpack weights
|
||||
unpack_weights = torch.bitwise_right_shift(
|
||||
torch.unsqueeze(qweights, 2).expand(-1, -1, 32 // bits),
|
||||
wf.unsqueeze(0),
|
||||
).to(torch.int16 if bits == 8 else torch.int8)
|
||||
unpack_weights = unpack_weights[:, :, order_map]
|
||||
unpack_weights = torch.bitwise_and(unpack_weights, (2**bits) - 1)
|
||||
# w : groups, group_size, n
|
||||
unpack_weights = unpack_weights.reshape(
|
||||
-1, group_size, unpack_weights.shape[1] * unpack_weights.shape[2]
|
||||
)
|
||||
|
||||
deq_weights = (unpack_weights - unpack_zeros) * scales.reshape(
|
||||
-1, 1, scales.shape[-1]
|
||||
)
|
||||
deq_weights = deq_weights.reshape(ic, oc)
|
||||
return deq_weights.to(output_type)
|
||||
|
||||
|
||||
def create_dynamic_map(signed=True, max_exponent_bits=7, total_bits=8):
|
||||
"""
|
||||
Creates the dynamic quantiztion map.
|
||||
|
||||
The dynamic data type is made up of a dynamic exponent and
|
||||
fraction. As the exponent increase from 0 to -7 the number
|
||||
of bits available for the fraction shrinks.
|
||||
|
||||
This is a generalization of the dynamic type where a certain
|
||||
number of the bits and be reserved for the linear quantization
|
||||
region (the fraction). n determines the maximum number of
|
||||
exponent bits.
|
||||
|
||||
For more details see
|
||||
(8-Bit Approximations for Parallelism in Deep Learning)[https://arxiv.org/abs/1511.04561]
|
||||
"""
|
||||
|
||||
data = []
|
||||
# these are additional items that come from the case
|
||||
# where all the exponent bits are zero and no
|
||||
# indicator bit is present
|
||||
non_sign_bits = total_bits - (1 if signed else 1)
|
||||
additional_items = 2 ** (non_sign_bits - max_exponent_bits) - 1
|
||||
for i in range(max_exponent_bits):
|
||||
fraction_items = int(
|
||||
2 ** (i + non_sign_bits - max_exponent_bits) + 1
|
||||
if signed
|
||||
else 2 ** (i + non_sign_bits - max_exponent_bits + 1) + 1,
|
||||
)
|
||||
boundaries = torch.linspace(0.1, 1, fraction_items)
|
||||
means = (boundaries[:-1] + boundaries[1:]) / 2.0
|
||||
data += ((10 ** (-(max_exponent_bits - 1) + i)) * means).tolist()
|
||||
if signed:
|
||||
data += (-(10 ** (-(max_exponent_bits - 1) + i)) * means).tolist()
|
||||
|
||||
if additional_items > 0:
|
||||
boundaries = torch.linspace(0.1, 1, additional_items + 1)
|
||||
means = (boundaries[:-1] + boundaries[1:]) / 2.0
|
||||
data += ((10 ** (-(max_exponent_bits - 1) + i)) * means).tolist()
|
||||
if signed:
|
||||
data += (-(10 ** (-(max_exponent_bits - 1) + i)) * means).tolist()
|
||||
|
||||
data.append(0)
|
||||
data.append(1.0)
|
||||
|
||||
assert len(data) == 2**total_bits
|
||||
|
||||
gap = 256 - len(data)
|
||||
for i in range(gap):
|
||||
data.append(0)
|
||||
|
||||
data.sort()
|
||||
return Tensor(data)
|
||||
|
||||
|
||||
def weight_quantize(
|
||||
A: torch.Tensor,
|
||||
absmax: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
blocksize: int,
|
||||
n: int,
|
||||
quant_dtype: str,
|
||||
code: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
A: (row ,col) torch.float16, torch.bfloat16, torch.float32
|
||||
absmax: (blocks) torch.float32
|
||||
blocks = n // blocksize, blocks += 1 if n % blocksize > 0 else 0
|
||||
quant_dtype: str
|
||||
目前可支持"int8"/"fp4"/"nf4"
|
||||
blocksize: int
|
||||
目前只支持4096, 2048, 1024, 512, 256, 128, 64
|
||||
n: int
|
||||
n = A.numel()
|
||||
out: (row ,col) torch.int8
|
||||
code: torch.float32
|
||||
the quantization map
|
||||
Returns:
|
||||
out: (row ,col) torch.int8
|
||||
|
||||
"""
|
||||
assert quant_dtype == "int8" or "fp4" or "nf4"
|
||||
if code is None and quant_dtype == "int8":
|
||||
code = create_dynamic_map().to(A.device)
|
||||
if isinstance(A, torch.Tensor) and not A.requires_grad:
|
||||
ops.infer.weight_quantize(A, absmax, out, blocksize, n, quant_dtype, code)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
47
ixformer_sdk/inference/functions/residual_bias.py
Normal file
47
ixformer_sdk/inference/functions/residual_bias.py
Normal file
@@ -0,0 +1,47 @@
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = ["residual_bias", "ref_residual_bias"]
|
||||
|
||||
|
||||
def ref_residual_bias(
|
||||
input: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
alpha: float = 1,
|
||||
):
|
||||
if bias is not None:
|
||||
output = residual.float() * alpha + input.float() + bias.float()
|
||||
else:
|
||||
output = residual.float() * alpha + input.float()
|
||||
|
||||
return output.to(residual.dtype)
|
||||
|
||||
|
||||
def residual_bias(
|
||||
input: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
bias: torch.Tensor = None,
|
||||
alpha: float = 1,
|
||||
output: torch.Tensor = None
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: [batch_count, seq_len, hidden_size] or [batch_tokens, hidden_size] torch.half
|
||||
residual: [batch_count, seq_len, hidden_size] or [batch_tokens, hidden_size] torch.half
|
||||
bias: [hidden_size] torch.half
|
||||
alpha: float
|
||||
Returns:
|
||||
output: [batch_count, seq_len, hidden_size] or [batch_tokens, hidden_size] torch.half
|
||||
"""
|
||||
if output is None:
|
||||
output = torch.empty_like(input)
|
||||
if alpha is None:
|
||||
alpha = 1
|
||||
if bias is not None:
|
||||
ops.train.add_residual_bias_forward(input, residual, bias, alpha, output)
|
||||
else:
|
||||
ops.train.add_residual_bias_forward(input, residual, alpha, output)
|
||||
return output
|
||||
134
ixformer_sdk/inference/functions/rms_norm.py
Normal file
134
ixformer_sdk/inference/functions/rms_norm.py
Normal file
@@ -0,0 +1,134 @@
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
from torch.nn import init
|
||||
|
||||
__all__ = ["ref_rms_norm", "rms_norm", "ref_residual_rms_norm", "residual_rms_norm"]
|
||||
|
||||
|
||||
def ref_rms_norm(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
eps: float = 1e-5,
|
||||
output: torch.Tensor = None,
|
||||
):
|
||||
dtype = input.dtype
|
||||
input = input.float()
|
||||
weight = weight.float()
|
||||
rms_out = input * torch.rsqrt(input.pow(2).mean(-1, keepdim=True) + eps)
|
||||
rms_out = rms_out * weight
|
||||
rms_out = rms_out.to(dtype)
|
||||
if output is not None:
|
||||
output.copy_(rms_out)
|
||||
else:
|
||||
output = rms_out
|
||||
return output
|
||||
|
||||
|
||||
def rms_norm(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
eps: float = 1e-5,
|
||||
output: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
This function is deprecated, please use residual_rms_norm.
|
||||
Args:
|
||||
input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
eps: float32
|
||||
Returns:
|
||||
output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
"""
|
||||
if output is None:
|
||||
output = torch.empty_like(input)
|
||||
|
||||
ops.infer.rms_norm(input, weight, output, None, eps)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def ref_residual_rms_norm(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
eps: float = 1e-5,
|
||||
residual_alpha: float = 1.0,
|
||||
residual: torch.Tensor = None,
|
||||
residual_bias: torch.Tensor = None,
|
||||
output: torch.Tensor = None,
|
||||
residual_output: torch.Tensor = None,
|
||||
is_post: bool = False,
|
||||
):
|
||||
dtype = input.dtype
|
||||
|
||||
if residual_bias is not None:
|
||||
input = input + residual_bias
|
||||
|
||||
if residual is not None:
|
||||
residual_output = torch.add(
|
||||
input, residual * residual_alpha, out=residual_output
|
||||
)
|
||||
input = input.float() + residual.float() * residual_alpha
|
||||
else:
|
||||
input = input.float()
|
||||
weight = weight.float()
|
||||
rms_out = input * torch.rsqrt(input.pow(2).mean(-1, keepdim=True) + eps)
|
||||
rms_out = rms_out * weight
|
||||
rms_out = rms_out.to(dtype)
|
||||
if output is not None:
|
||||
output.copy_(rms_out)
|
||||
else:
|
||||
output = rms_out
|
||||
|
||||
if is_post and residual_output is not None:
|
||||
residual_output = output
|
||||
|
||||
return output, residual_output
|
||||
|
||||
|
||||
def residual_rms_norm(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
eps: float = 1e-5,
|
||||
residual_alpha: float = 1.0,
|
||||
residual: torch.Tensor = None,
|
||||
residual_bias: torch.Tensor = None,
|
||||
output: torch.Tensor = None,
|
||||
residual_output: torch.Tensor = None,
|
||||
is_post: bool = False,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
eps: float32
|
||||
residual_alpha: float32
|
||||
residual: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
residual_bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
is_post: bool
|
||||
Returns:
|
||||
output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 If set to None, an inplace operation will be performed on input.
|
||||
residual_output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 If set to None, an inplace operation will be performed on residual.
|
||||
"""
|
||||
if residual is None:
|
||||
if output is None:
|
||||
output = torch.empty_like(input)
|
||||
|
||||
ops.infer.rms_norm(input, weight, output, residual_bias, eps)
|
||||
else:
|
||||
ops.infer.residual_rms_norm(
|
||||
input,
|
||||
residual,
|
||||
weight,
|
||||
output,
|
||||
residual_output,
|
||||
residual_bias,
|
||||
residual_alpha,
|
||||
eps,
|
||||
is_post,
|
||||
)
|
||||
residual_output = residual_output if residual_output is not None else residual
|
||||
output = output if output is not None else input
|
||||
|
||||
return output, residual_output
|
||||
120
ixformer_sdk/inference/functions/scaled_dot_product_attention.py
Normal file
120
ixformer_sdk/inference/functions/scaled_dot_product_attention.py
Normal file
@@ -0,0 +1,120 @@
|
||||
import math
|
||||
from typing import List, Union
|
||||
|
||||
import torch
|
||||
|
||||
from .flash_attn import ixinfer_flash_attn_pad
|
||||
|
||||
__all__ = ["scaled_dot_product_attention", "ref_scaled_dot_product_attention"]
|
||||
|
||||
|
||||
def ref_scaled_dot_product_attention(
|
||||
query: "torch.Tensor",
|
||||
key: "torch.Tensor",
|
||||
value: "torch.Tensor",
|
||||
attn_mask=None,
|
||||
dropout_p=0.0,
|
||||
is_causal=False,
|
||||
):
|
||||
assert len(query.shape) >= 3
|
||||
assert len(key.shape) >= 3
|
||||
assert len(value.shape) >= 3
|
||||
batch_size = query.shape[0]
|
||||
L = query.shape[-2]
|
||||
S = key.shape[-2]
|
||||
|
||||
if is_causal and attn_mask is not None:
|
||||
raise RuntimeError()
|
||||
|
||||
if attn_mask is None and is_causal is False:
|
||||
attn_mask = torch.ones([batch_size, 1, 1, S]).bool().to(query.device)
|
||||
|
||||
if is_causal:
|
||||
attn_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0).to(query.device)
|
||||
if attn_mask.dtype == torch.bool:
|
||||
attn_mask = (
|
||||
torch.zeros_like(attn_mask).to(query.dtype).masked_fill(~attn_mask, -10000)
|
||||
)
|
||||
# attn_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0) if is_causal else attn_mask
|
||||
# attn_mask = attn_mask.masked_fill(not attn_mask, -float('inf')) if attn_mask.dtype==torch.bool else attn_mask
|
||||
attn_weight = torch.softmax(
|
||||
(query @ key.transpose(-2, -1) / math.sqrt(query.size(-1))) + attn_mask, dim=-1
|
||||
)
|
||||
attn_weight = torch.dropout(attn_weight, dropout_p, True)
|
||||
return attn_weight.to(query.dtype) @ value
|
||||
|
||||
|
||||
def scaled_dot_product_attention(
|
||||
query: "torch.Tensor",
|
||||
key: "torch.Tensor",
|
||||
value: "torch.Tensor",
|
||||
attn_mask=None,
|
||||
dropout_p=0.0,
|
||||
is_causal=False,
|
||||
):
|
||||
# 1. pytorch版本 query,key的head_dim相等， value可以不相等。但是我们实现的版本需要相等
|
||||
# 2. pytorch版本的 L, S可以不相等， 但是我们必须相等，并且为64的倍数
|
||||
# 3. pytorch支持 加上 mask
|
||||
# 4. mask: 0 表示padding[与后端相反，所以注意如果用传进来的，取反转int；如果是自己生成，则直接生成后端的mask]
|
||||
|
||||
"""
|
||||
Args:
|
||||
query: (N, ..., L, E) torch.float16, torch.bfloat16
|
||||
key: (N, ..., S, E) torch.float16, torch.bfloat16
|
||||
value: (N, ..., S, E) torch.float16, torch.bfloat16
|
||||
attn_mask: (N, ..., L, S) bool, torch.float32
|
||||
dropout_p: float32
|
||||
Dropout probability; if greater than 0.0, dropout is applied
|
||||
is_causal: bool
|
||||
If true, assumes causal attention masking and errors if both attn_mask and is_causal are set.
|
||||
Returns:
|
||||
Tensor: (N, ..., L, E) torch.float16, torch.bfloat16
|
||||
"""
|
||||
|
||||
assert len(query.shape) >= 4, "len(query.shape) <4"
|
||||
assert len(key.shape) >= 4, "len(key.shape) <4"
|
||||
assert len(value.shape) >= 4, "len(value.shape) <4"
|
||||
query_shape = list(query.shape)
|
||||
key_shape = list(key.shape)
|
||||
value_shape = list(value.shape)
|
||||
|
||||
batch_size = query_shape[0]
|
||||
q_seq_len = query_shape[-2]
|
||||
kv_seq_len = key_shape[-2]
|
||||
|
||||
if len(query_shape) > 4:
|
||||
query.view(batch_size, -1, q_seq_len, query_shape[-1])
|
||||
if len(key_shape) > 4:
|
||||
key.view(batch_size, -1, kv_seq_len, key_shape[-1])
|
||||
if len(value_shape) > 4:
|
||||
value.view(batch_size, -1, kv_seq_len, value_shape[-1])
|
||||
|
||||
training = query.requires_grad
|
||||
atten_scale = 1.0 / (query.size(-1) ** 0.5)
|
||||
|
||||
if is_causal and attn_mask is not None: # 两者必须有有一个
|
||||
raise RuntimeError()
|
||||
|
||||
head_num = query.size(1)
|
||||
# 注意，training,inference都支持mask广播；
|
||||
if training:
|
||||
raise NotImplementedError("not support training!")
|
||||
|
||||
else: # inference支持mask广播；
|
||||
if attn_mask is None and is_causal is False: # mask 全0
|
||||
attn_mask = None
|
||||
# attn_mask = (
|
||||
# torch.zeros([batch_size, 1, 1, kv_seq_len]).int().to(query.device)
|
||||
# ) # 底层代码1代表mask，0代表保留
|
||||
elif is_causal: # mask 下三角是0，上三角是1
|
||||
attn_mask = (
|
||||
torch.ones([batch_size, 1, q_seq_len, kv_seq_len], dtype=torch.int)
|
||||
.triu(diagonal=1)
|
||||
.to(query.device)
|
||||
) # 上三角是1，下三角和对角线是0
|
||||
elif attn_mask is not None: # 外部传进来的，取反，转int
|
||||
assert attn_mask.dtype == torch.bool or attn_mask.dtype == torch.float
|
||||
assert attn_mask.dim() == 4 # 必须是4维
|
||||
if attn_mask.dtype == torch.bool:
|
||||
attn_mask = (~attn_mask).int() # 取非操作，然后转int
|
||||
return ixinfer_flash_attn_pad(query, key, value, attn_mask, None, atten_scale)
|
||||
510
ixformer_sdk/inference/functions/smoothquant.py
Normal file
510
ixformer_sdk/inference/functions/smoothquant.py
Normal file
@@ -0,0 +1,510 @@
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
import torch.nn.functional as NNF
|
||||
|
||||
__all__ = [
|
||||
"ref_dynamic_scaled_quant_dynamic_int8",
|
||||
"dynamic_scaled_quant_dynamic_int8",
|
||||
"dynamic_scaled_quant_smoothquant",
|
||||
"ref_silu_and_mul_smoothquant",
|
||||
"silu_and_mul_smoothquant",
|
||||
"ref_residual_rms_norm_dynamic_int8",
|
||||
"residual_rms_norm_dynamic_int8",
|
||||
"ref_residual_layer_norm_dynamic_int8",
|
||||
"residual_layer_norm_dynamic_int8",
|
||||
"ref_layer_norm_2sb_smoothquant",
|
||||
"layer_norm_2sb_smoothquant",
|
||||
"ref_residual_layer_norm_2sb_smoothquant",
|
||||
"residual_layer_norm_2sb_smoothquant",
|
||||
]
|
||||
|
||||
|
||||
def ref_dynamic_scaled_quant_dynamic_int8(
|
||||
input: torch.Tensor,
|
||||
smooth_scales: torch.Tensor = None,
|
||||
i8_output: torch.Tensor = None,
|
||||
output_scales: torch.Tensor = None,
|
||||
):
|
||||
if i8_output is None:
|
||||
i8_output = torch.empty(input.shape, dtype=torch.int8, device=input.device)
|
||||
if output_scales is None:
|
||||
output_scales = torch.empty(
|
||||
input.shape[:-1], dtype=torch.float32, device=input.device
|
||||
)
|
||||
|
||||
scales_shape = input.shape[:-1]
|
||||
output = input.float()
|
||||
if smooth_scales is not None:
|
||||
output *= smooth_scales.view(1, -1)
|
||||
amax_, _ = torch.max(torch.abs(output), dim=-1, keepdim=True)
|
||||
scales = amax_ / 127.0
|
||||
output = output / scales
|
||||
output = torch.clamp(torch.round(output), -127, 127).to(torch.int8)
|
||||
|
||||
if i8_output is not None:
|
||||
i8_output.copy_(output)
|
||||
output = i8_output
|
||||
if output_scales is not None:
|
||||
output_scales.view(-1).copy_(scales.view(-1))
|
||||
scales = output_scales
|
||||
|
||||
return output, scales.view(scales_shape)
|
||||
|
||||
|
||||
def dynamic_scaled_quant_dynamic_int8(
|
||||
input: torch.Tensor,
|
||||
smooth_scales: torch.Tensor = None,
|
||||
i8_output: torch.Tensor = None,
|
||||
output_scales: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (..., k) torch.float16,torch.bfloat16
|
||||
smooth_scales: (k) torch.float16,torch.bfloat16
|
||||
if smooth_scales is None, api is dynamic-per-token quantization.
|
||||
Returns:
|
||||
i8_output: (..., k) torch.int8
|
||||
output_scales: (...) torch.float32
|
||||
"""
|
||||
if i8_output is None:
|
||||
i8_output = torch.empty(input.shape, dtype=torch.int8, device=input.device)
|
||||
if output_scales is None:
|
||||
output_scales = torch.empty(
|
||||
input.shape[:-1], dtype=torch.float32, device=input.device
|
||||
)
|
||||
hidden_size = input.shape[-1]
|
||||
|
||||
if smooth_scales is None:
|
||||
ops.infer.scaled_int8_quant(i8_output, input, output_scales, 1)
|
||||
return i8_output, output_scales
|
||||
|
||||
ops.infer.dynamic_scaled_quant_smoothquant(
|
||||
input.view(-1, hidden_size),
|
||||
smooth_scales,
|
||||
i8_output.view(-1, hidden_size),
|
||||
output_scales,
|
||||
)
|
||||
|
||||
return i8_output, output_scales
|
||||
|
||||
|
||||
# For backward compatibility
|
||||
dynamic_scaled_quant_smoothquant = dynamic_scaled_quant_dynamic_int8
|
||||
|
||||
|
||||
def ref_silu_and_mul_smoothquant(
|
||||
input, smooth_scales, i8_output=None, output_scales=None
|
||||
):
|
||||
x1, x2 = input.chunk(chunks=2, dim=-1)
|
||||
x = NNF.silu(x1) * x2
|
||||
|
||||
return ref_dynamic_scaled_quant_dynamic_int8(
|
||||
x, smooth_scales, i8_output, output_scales
|
||||
)
|
||||
|
||||
|
||||
def silu_and_mul_smoothquant(input, smooth_scales, i8_output=None, output_scales=None):
|
||||
"""
|
||||
Args:
|
||||
input: (..., 2*k) torch.float16,torch.bfloat16
|
||||
smooth_scales: (k) torch.float16,torch.bfloat16
|
||||
if smooth_scales is None, api is dynamic-per-token quantization.
|
||||
Returns:
|
||||
i8_output: (..., k) torch.int8
|
||||
output_scales: (...) torch.float32
|
||||
"""
|
||||
if i8_output is None:
|
||||
output_shape = input.shape[:-1] + (input.shape[-1] // 2,)
|
||||
i8_output = torch.empty(output_shape, dtype=torch.int8, device=input.device)
|
||||
if output_scales is None:
|
||||
output_scales = torch.empty(
|
||||
input.shape[:-1], dtype=torch.float32, device=input.device
|
||||
)
|
||||
|
||||
ops.infer.silu_and_mul_smoothquant(i8_output, input, smooth_scales, output_scales)
|
||||
|
||||
return i8_output, output_scales
|
||||
|
||||
|
||||
def ref_residual_rms_norm_dynamic_int8(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
residual: torch.Tensor = None,
|
||||
residual_bias: torch.Tensor = None,
|
||||
eps: float = 1e-5,
|
||||
smooth_scales: torch.Tensor = None,
|
||||
output: torch.Tensor = None,
|
||||
residual_output: torch.Tensor = None,
|
||||
output_scales: torch.Tensor = None,
|
||||
is_post: bool = False,
|
||||
):
|
||||
dtype = input.dtype
|
||||
if output is None:
|
||||
output = torch.empty(input.shape, dtype=torch.int8, device=input.device)
|
||||
if output_scales is None:
|
||||
output_scales = torch.empty(
|
||||
input.shape[:-1], dtype=torch.float32, device=input.device
|
||||
)
|
||||
if residual_bias is not None:
|
||||
input = input + residual_bias
|
||||
|
||||
if residual is not None:
|
||||
residual_output = torch.add(input, residual, out=residual_output)
|
||||
input = residual_output
|
||||
input = input.float()
|
||||
weight = weight.float()
|
||||
rms_output = input * torch.rsqrt(input.pow(2).mean(-1, keepdim=True) + eps)
|
||||
rms_output = (rms_output * weight).to(dtype)
|
||||
|
||||
if residual is not None and is_post:
|
||||
residual_output.copy_(rms_output)
|
||||
output, output_scales = ref_dynamic_scaled_quant_dynamic_int8(
|
||||
rms_output, smooth_scales, output, output_scales.view(-1)
|
||||
)
|
||||
|
||||
return output, residual_output, output_scales
|
||||
|
||||
|
||||
def residual_rms_norm_dynamic_int8(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
residual: torch.Tensor = None,
|
||||
residual_bias: torch.Tensor = None,
|
||||
eps: float = 1e-5,
|
||||
smooth_scales: torch.Tensor = None,
|
||||
output: torch.Tensor = None,
|
||||
residual_output: torch.Tensor = None,
|
||||
output_scales: torch.Tensor = None,
|
||||
is_post: bool = False,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
residual: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
residual_bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
eps: float32
|
||||
smooth_scales: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
is_post: bool
|
||||
Returns:
|
||||
output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
residual_output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 If set to None, an inplace operation will be performed on residual.
|
||||
output_scales: (...) torch.float16, torch.bfloat16, torch.float32
|
||||
"""
|
||||
if output is None:
|
||||
output = torch.empty(input.shape, dtype=torch.int8, device=input.device)
|
||||
if output_scales is None:
|
||||
output_scales = torch.empty(
|
||||
input.shape[:-1], dtype=torch.float32, device=input.device
|
||||
)
|
||||
|
||||
if residual is None:
|
||||
ops.infer.rmsnorm_dynamic_int8(
|
||||
input, weight, output, output_scales, smooth_scales, residual_bias, eps
|
||||
)
|
||||
else:
|
||||
ops.infer.residual_rmsnorm_dynamic_int8(
|
||||
input,
|
||||
residual,
|
||||
weight,
|
||||
output,
|
||||
output_scales,
|
||||
smooth_scales,
|
||||
residual_output,
|
||||
residual_bias,
|
||||
eps,
|
||||
is_post,
|
||||
)
|
||||
residual_output = residual if residual_output is None else residual_output
|
||||
|
||||
return output, residual_output, output_scales
|
||||
|
||||
|
||||
def ref_residual_layer_norm_dynamic_int8(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
residual: torch.Tensor = None,
|
||||
residual_bias: torch.Tensor = None,
|
||||
eps: float = 1e-5,
|
||||
smooth_scales: torch.Tensor = None,
|
||||
output: torch.Tensor = None,
|
||||
residual_output: torch.Tensor = None,
|
||||
output_scales: torch.Tensor = None,
|
||||
):
|
||||
normalized_shape = [weight.size(-1)]
|
||||
|
||||
if output is None:
|
||||
output = torch.empty(input.shape, dtype=torch.int8, device=input.device)
|
||||
if output_scales is None:
|
||||
output_scales = torch.empty(
|
||||
input.shape[:-1], dtype=torch.float32, device=input.device
|
||||
)
|
||||
|
||||
if residual_bias is not None:
|
||||
input = input + residual_bias
|
||||
|
||||
if residual is not None:
|
||||
residual_output = torch.add(input, residual, out=residual_output)
|
||||
input = residual_output
|
||||
|
||||
norm_output = torch.nn.functional.layer_norm(
|
||||
input, normalized_shape, weight, bias, eps=eps
|
||||
)
|
||||
|
||||
output, output_scales = ref_dynamic_scaled_quant_dynamic_int8(
|
||||
norm_output, smooth_scales, output, output_scales.view(-1)
|
||||
)
|
||||
|
||||
return output, residual_output, output_scales
|
||||
|
||||
|
||||
def residual_layer_norm_dynamic_int8(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
bias: torch.Tensor,
|
||||
residual: torch.Tensor = None,
|
||||
residual_bias: torch.Tensor = None,
|
||||
eps: float = 1e-5,
|
||||
smooth_scales: torch.Tensor = None,
|
||||
output: torch.Tensor = None,
|
||||
residual_output: torch.Tensor = None,
|
||||
output_scales: torch.Tensor = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
weight: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
residual: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
residual_bias: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
eps: float32
|
||||
smooth_scales: (hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
Returns:
|
||||
output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32
|
||||
residual_output: (..., hidden_size) torch.float16, torch.bfloat16, torch.float32 If set to None, an inplace operation will be performed on residual.
|
||||
output_scales: (...) torch.float16, torch.bfloat16, torch.float32
|
||||
"""
|
||||
if output is None:
|
||||
output = torch.empty(input.shape, dtype=torch.int8, device=input.device)
|
||||
if output_scales is None:
|
||||
output_scales = torch.empty(
|
||||
input.shape[:-1], dtype=torch.float32, device=input.device
|
||||
)
|
||||
|
||||
if residual is None:
|
||||
ops.infer.layer_norm_dynamic_int8(
|
||||
input,
|
||||
weight,
|
||||
bias,
|
||||
output,
|
||||
output_scales,
|
||||
smooth_scales,
|
||||
residual_bias,
|
||||
eps,
|
||||
)
|
||||
else:
|
||||
ops.infer.residual_layer_norm_dynamic_int8(
|
||||
input,
|
||||
residual,
|
||||
weight,
|
||||
bias,
|
||||
output,
|
||||
output_scales,
|
||||
smooth_scales,
|
||||
residual_output,
|
||||
residual_bias,
|
||||
eps,
|
||||
)
|
||||
residual_output = residual_output if residual_output is not None else residual
|
||||
|
||||
return output, residual_output, output_scales
|
||||
|
||||
|
||||
def ref_layer_norm_2sb_smoothquant(
|
||||
input,
|
||||
weight1,
|
||||
bias1,
|
||||
smooth_scales1,
|
||||
weight2,
|
||||
bias2,
|
||||
smooth_scales2,
|
||||
i8_output1=None,
|
||||
output_scales1=None,
|
||||
i8_output2=None,
|
||||
output_scales2=None,
|
||||
eps=1e-5,
|
||||
):
|
||||
input1 = torch.nn.functional.layer_norm(
|
||||
input, [weight1.shape[-1]], weight1, bias1, eps=eps
|
||||
)
|
||||
input2 = torch.nn.functional.layer_norm(
|
||||
input, [weight2.shape[-1]], weight2, bias2, eps=eps
|
||||
)
|
||||
|
||||
i8_output1, output_scales1 = ref_dynamic_scaled_quant_dynamic_int8(
|
||||
input1, smooth_scales1, i8_output1, output_scales1
|
||||
)
|
||||
i8_output2, output_scales2 = ref_dynamic_scaled_quant_dynamic_int8(
|
||||
input2, smooth_scales2, i8_output2, output_scales2
|
||||
)
|
||||
|
||||
return i8_output1, output_scales1, i8_output2, output_scales2
|
||||
|
||||
|
||||
def layer_norm_2sb_smoothquant(
|
||||
input,
|
||||
weight1,
|
||||
bias1,
|
||||
smooth_scales1,
|
||||
weight2,
|
||||
bias2,
|
||||
smooth_scales2,
|
||||
output1=None,
|
||||
output_scales1=None,
|
||||
output2=None,
|
||||
output_scales2=None,
|
||||
eps=1e-5,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (..., hidden_size) torch.float16, torch.bfloat16
|
||||
weight1: (hidden_size) torch.float16, torch.bfloat16
|
||||
bias1: (hidden_size) torch.float16, torch.bfloat16
|
||||
smooth_scales1: (hidden_size) torch.float16, torch.bfloat16
|
||||
weight2: (hidden_size) torch.float16, torch.bfloat16
|
||||
bias2: (hidden_size) torch.float16, torch.bfloat16
|
||||
smooth_scales2: (hidden_size) torch.float16, torch.bfloat16
|
||||
eps: float32
|
||||
Returns:
|
||||
output1: (..., hidden_size) torch.float16, torch.bfloat16
|
||||
output_scales1: (...) torch.float16, torch.bfloat16
|
||||
output2: (..., hidden_size) torch.float16, torch.bfloat16
|
||||
output_scales2: (...) torch.float16, torch.bfloat16
|
||||
"""
|
||||
if output1 is None:
|
||||
output1 = torch.empty(input.shape, dtype=torch.int8, device=input.device)
|
||||
if output_scales1 is None:
|
||||
output_scales1 = torch.empty(
|
||||
input.shape[:-1], dtype=torch.float32, device=input.device
|
||||
)
|
||||
if output2 is None:
|
||||
output2 = torch.empty(input.shape, dtype=torch.int8, device=input.device)
|
||||
if output_scales2 is None:
|
||||
output_scales2 = torch.empty(
|
||||
input.shape[:-1], dtype=torch.float32, device=input.device
|
||||
)
|
||||
|
||||
ops.infer.layer_norm_2sb_smoothquant(
|
||||
input,
|
||||
weight1,
|
||||
bias1,
|
||||
smooth_scales1,
|
||||
weight2,
|
||||
bias2,
|
||||
smooth_scales2,
|
||||
output1,
|
||||
output_scales1,
|
||||
output2,
|
||||
output_scales2,
|
||||
eps,
|
||||
)
|
||||
|
||||
return output1, output_scales1, output2, output_scales2
|
||||
|
||||
|
||||
def ref_residual_layer_norm_2sb_smoothquant(
|
||||
input,
|
||||
residual,
|
||||
weight1,
|
||||
bias1,
|
||||
smooth_scales1,
|
||||
weight2,
|
||||
bias2,
|
||||
smooth_scales2,
|
||||
i8_output1=None,
|
||||
output_scales1=None,
|
||||
i8_output2=None,
|
||||
output_scales2=None,
|
||||
eps=1e-5,
|
||||
):
|
||||
residual_out = input + residual
|
||||
input1 = torch.nn.functional.layer_norm(
|
||||
residual_out, [weight1.shape[-1]], weight1, bias1, eps=eps
|
||||
)
|
||||
input2 = torch.nn.functional.layer_norm(
|
||||
residual_out, [weight2.shape[-1]], weight2, bias2, eps=eps
|
||||
)
|
||||
|
||||
i8_output1, output_scales1 = ref_dynamic_scaled_quant_dynamic_int8(
|
||||
input1, smooth_scales1, i8_output1, output_scales1
|
||||
)
|
||||
i8_output2, output_scales2 = ref_dynamic_scaled_quant_dynamic_int8(
|
||||
input2, smooth_scales2, i8_output2, output_scales2
|
||||
)
|
||||
|
||||
return residual_out, i8_output1, output_scales1, i8_output2, output_scales2
|
||||
|
||||
|
||||
def residual_layer_norm_2sb_smoothquant(
|
||||
input,
|
||||
residual,
|
||||
weight1,
|
||||
bias1,
|
||||
smooth_scales1,
|
||||
weight2,
|
||||
bias2,
|
||||
smooth_scales2,
|
||||
output1=None,
|
||||
output_scales1=None,
|
||||
output2=None,
|
||||
output_scales2=None,
|
||||
eps=1e-5,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (..., hidden_size) torch.float16, torch.bfloat16
|
||||
residual: (..., hidden_size) torch.float16, torch.bfloat16
|
||||
weight1: (hidden_size) torch.float16, torch.bfloat16
|
||||
bias1: (hidden_size) torch.float16, torch.bfloat16
|
||||
smooth_scales1: (hidden_size) torch.float16, torch.bfloat16
|
||||
weight2: (hidden_size) torch.float16, torch.bfloat16
|
||||
bias2: (hidden_size) torch.float16, torch.bfloat16
|
||||
smooth_scales2: (hidden_size) torch.float16, torch.bfloat16
|
||||
eps: float32
|
||||
Returns:
|
||||
output1: (..., hidden_size) torch.float16, torch.bfloat16
|
||||
output_scales1: (...) torch.float16, torch.bfloat16
|
||||
output2: (..., hidden_size) torch.float16, torch.bfloat16
|
||||
output_scales2: (...) torch.float16, torch.bfloat16
|
||||
"""
|
||||
if output1 is None:
|
||||
output1 = torch.empty(input.shape, dtype=torch.int8, device=input.device)
|
||||
if output_scales1 is None:
|
||||
output_scales1 = torch.empty(
|
||||
input.shape[:-1], dtype=torch.float32, device=input.device
|
||||
)
|
||||
if output2 is None:
|
||||
output2 = torch.empty(input.shape, dtype=torch.int8, device=input.device)
|
||||
if output_scales2 is None:
|
||||
output_scales2 = torch.empty(
|
||||
input.shape[:-1], dtype=torch.float32, device=input.device
|
||||
)
|
||||
|
||||
ops.infer.residual_layer_norm_2sb_smoothquant(
|
||||
input,
|
||||
residual,
|
||||
weight1,
|
||||
bias1,
|
||||
smooth_scales1,
|
||||
weight2,
|
||||
bias2,
|
||||
smooth_scales2,
|
||||
output1,
|
||||
output_scales1,
|
||||
output2,
|
||||
output_scales2,
|
||||
eps,
|
||||
)
|
||||
|
||||
return residual, output1, output_scales1, output2, output_scales2
|
||||
33
ixformer_sdk/inference/functions/softmax.py
Normal file
33
ixformer_sdk/inference/functions/softmax.py
Normal file
@@ -0,0 +1,33 @@
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
from torch.autograd.function import Function, FunctionCtx
|
||||
|
||||
__all__ = ["softmax", "ref_softmax"]
|
||||
|
||||
|
||||
def ref_softmax(input: torch.Tensor, dim: int = None, _stacklevel: int = 3, dtype=None):
|
||||
out = torch.nn.functional.softmax(
|
||||
input, dim=dim, _stacklevel=_stacklevel, dtype=dtype
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def softmax(input: torch.Tensor, dim=None, _stacklevel=3, dtype=None):
|
||||
|
||||
"""
|
||||
Args:
|
||||
input: (...) torch.float16
|
||||
dim: int
|
||||
要进行softmax的维度,目前只支持最后一维， dim==-1 or dim == input.dim()-1
|
||||
_stacklevel: int
|
||||
这个参数只是为了与pytorch中对齐。 stacklevel is used in python to indicate warning mechanism how far up the stack it has to go to find the line that called the function which issued the warning.
|
||||
dtype: torch.float16
|
||||
Returns:
|
||||
Tensor: (...) torch.float16
|
||||
"""
|
||||
output = torch.empty_like(input)
|
||||
ops.infer.softmax(input, output, dim)
|
||||
output = output.to(dtype)
|
||||
return output
|
||||
75
ixformer_sdk/inference/functions/store_kv_cache.py
Normal file
75
ixformer_sdk/inference/functions/store_kv_cache.py
Normal file
@@ -0,0 +1,75 @@
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"store_kv_cache",
|
||||
"ref_store_kv_cache",
|
||||
]
|
||||
|
||||
|
||||
def ref_store_kv_cache(
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
v_cache: torch.Tensor,
|
||||
cache_batch_idx: torch.Tensor,
|
||||
cache_seqlens: torch.Tensor,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
k: (batch_size, seqlen_new, head_num, head_dim) torch.float16, torch.bfloat16
|
||||
v: (batch_size, seqlen_new, head_num, head_dim) torch.float16, torch.bfloat16
|
||||
k_cache: (batch_size_cache, seqlen_cache, head_num, head_dim) torch.float16, torch.bfloat16
|
||||
v_cache: (batch_size_cache, seqlen_cache, head_num, head_dim) torch.float16, torch.bfloat16
|
||||
cache_batch_idx: (batch_size,) torch.int32
|
||||
The indices used to index into the KV cache.
|
||||
cache_seqlens: (batch_size,) torch.int32
|
||||
The sequence lengths of the KV cache.
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# 等价实现
|
||||
# concatenate k with k_cache, starting at the indices specified by cache_seqlens.
|
||||
seqlen_new = k.size(1)
|
||||
for kv_batch_idx, cache_kv_batch_idx in enumerate(cache_batch_idx):
|
||||
cache_len = cache_seqlens[kv_batch_idx]
|
||||
cache_start_idx = cache_len
|
||||
cache_end_idx = cache_len + seqlen_new
|
||||
|
||||
k_cache[cache_kv_batch_idx, cache_start_idx:cache_end_idx] = k[kv_batch_idx]
|
||||
v_cache[cache_kv_batch_idx, cache_start_idx:cache_end_idx] = v[kv_batch_idx]
|
||||
|
||||
|
||||
def store_kv_cache(
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
v_cache: torch.Tensor,
|
||||
cache_batch_idx: torch.Tensor,
|
||||
cache_seqlens: torch.Tensor,
|
||||
):
|
||||
"""
|
||||
Currently, only head_dim%2==0 is supported.
|
||||
Args:
|
||||
k: (batch_size, seqlen_new, head_num, head_dim) torch.float16, torch.bfloat16
|
||||
v: (batch_size, seqlen_new, head_num, head_dim) torch.float16, torch.bfloat16
|
||||
k_cache: (batch_size_cache, seqlen_cache, head_num, head_dim) torch.float16, torch.bfloat16
|
||||
v_cache: (batch_size_cache, seqlen_cache, head_num, head_dim) torch.float16, torch.bfloat16
|
||||
cache_batch_idx: (batch_size,) torch.int32
|
||||
The indices used to index into the KV cache.
|
||||
cache_seqlens: (batch_size,) torch.int32
|
||||
The sequence lengths of the KV cache
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
head_dim = k.size(-1)
|
||||
assert head_dim % 2 == 0, "Currently, only head_dim%2==0 is supported."
|
||||
|
||||
ops.infer.store_kv_cache(
|
||||
k,
|
||||
v,
|
||||
k_cache,
|
||||
v_cache,
|
||||
cache_batch_idx,
|
||||
cache_seqlens,
|
||||
)
|
||||
97
ixformer_sdk/inference/functions/t5.py
Normal file
97
ixformer_sdk/inference/functions/t5.py
Normal file
@@ -0,0 +1,97 @@
|
||||
from typing import Union
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"t5_split_qkv",
|
||||
"t5_split_qkv_update_kv_cache",
|
||||
"ref_t5_split_qkv_update_kv_cache",
|
||||
"ref_t5_split_qkv",
|
||||
]
|
||||
|
||||
|
||||
def reshape_query(query, head_num, head_dim):
|
||||
batch_size, seq_len, _ = query.shape
|
||||
query = query.view(batch_size, seq_len, head_num, head_dim)
|
||||
query = query.transpose(1, 2)
|
||||
return query
|
||||
|
||||
|
||||
def ref_t5_split_qkv(qkv: "torch.Tensor", head_num: int, head_dim: int):
|
||||
assert qkv.size(-1) == head_dim * head_num * 3
|
||||
batch_size, seq_len, _ = qkv.shape
|
||||
q, k, v = torch.chunk(qkv, 3, dim=-1)
|
||||
q = reshape_query(q, head_num, head_dim)
|
||||
k = reshape_query(k, head_num, head_dim)
|
||||
v = reshape_query(v, head_num, head_dim)
|
||||
return q, k, v
|
||||
|
||||
|
||||
def t5_split_qkv(qkv: "torch.Tensor", head_num: int, head_dim: int):
|
||||
|
||||
"""
|
||||
Args:
|
||||
qkv: (batch_size, seq_len, head_dim * head_num * 3) torch.half, torch.bfloat16
|
||||
head_num: int
|
||||
head_dim: int
|
||||
Returns:
|
||||
q: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16
|
||||
k: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16
|
||||
v: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16
|
||||
"""
|
||||
batch_size, seq_len, _ = qkv.shape
|
||||
q = qkv.new_empty([batch_size, head_num, seq_len, head_dim])
|
||||
k = qkv.new_empty([batch_size, head_num, seq_len, head_dim])
|
||||
v = qkv.new_empty([batch_size, head_num, seq_len, head_dim])
|
||||
ops.infer.t5_split_qkv(qkv, q, k, v, head_num, head_dim)
|
||||
return q, k, v
|
||||
|
||||
|
||||
def ref_t5_split_qkv_update_kv_cache(
|
||||
qkv: "torch.Tensor",
|
||||
past_key: "torch.Tensor",
|
||||
past_value: "torch.Tensor",
|
||||
head_num: int,
|
||||
head_dim: int,
|
||||
):
|
||||
assert qkv.size(-1) == head_dim * head_num * 3
|
||||
batch_size, seq_len, _ = qkv.shape
|
||||
q, k, v = torch.chunk(qkv, 3, dim=-1)
|
||||
q = reshape_query(q, head_num, head_dim)
|
||||
k = reshape_query(k, head_num, head_dim)
|
||||
v = reshape_query(v, head_num, head_dim)
|
||||
k = torch.cat([past_key, k], dim=2)
|
||||
v = torch.cat([past_value, v], dim=2)
|
||||
return q, k, v
|
||||
|
||||
|
||||
def t5_split_qkv_update_kv_cache(
|
||||
qkv: "torch.Tensor",
|
||||
past_key: "torch.Tensor",
|
||||
past_value: "torch.Tensor",
|
||||
head_num: int,
|
||||
head_dim: int,
|
||||
):
|
||||
|
||||
"""
|
||||
Args:
|
||||
qkv: (batch_size, 1 , head_dim * head_num * 3) torch.half, torch.bfloat16
|
||||
past_key: (batch_size, head_num, seq_len - 1, head_dim) torch.half, torch.bfloat16
|
||||
past_value: (batch_size, head_num, seq_len - 1, head_dim) torch.half, torch.bfloat16
|
||||
head_num: int
|
||||
head_dim: int
|
||||
Returns:
|
||||
q: (batch_size, head_num, 1, head_dim) torch.half, torch.bfloat16
|
||||
k: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16
|
||||
v: (batch_size, head_num, seq_len, head_dim) torch.half, torch.bfloat16
|
||||
"""
|
||||
batch_size, _, past_seq_len, _ = list(past_key.shape)
|
||||
seq_len = past_seq_len + 1
|
||||
q = qkv.new_empty([batch_size, head_num, 1, head_dim])
|
||||
k = qkv.new_empty([batch_size, head_num, seq_len, head_dim])
|
||||
v = qkv.new_empty([batch_size, head_num, seq_len, head_dim])
|
||||
ops.infer.t5_split_qkv_update_kv_cache(
|
||||
qkv, past_key, past_value, q, k, v, head_num, head_dim
|
||||
)
|
||||
return q, k, v
|
||||
617
ixformer_sdk/inference/functions/tgi.py
Normal file
617
ixformer_sdk/inference/functions/tgi.py
Normal file
@@ -0,0 +1,617 @@
|
||||
import math
|
||||
from typing import List, Optional
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"tgi_apply_rotary_emb_torch",
|
||||
"tgi_apply_rotary",
|
||||
"tgi_gather_prefill_logprobs",
|
||||
"ref_paged_attention_v1",
|
||||
"ref_paged_attention_v3",
|
||||
"get_alibi_slopes",
|
||||
"paged_attention_v1",
|
||||
"reshape_and_cache_v1",
|
||||
"paged_attention_v7",
|
||||
"reshape_and_cache",
|
||||
"paged_attention_v3",
|
||||
"ref_reshape_and_cache_v3",
|
||||
"reshape_and_cache_v3",
|
||||
]
|
||||
|
||||
|
||||
def get_alibi_slopes(total_num_heads: int) -> torch.Tensor:
|
||||
closest_power_of_2 = 2 ** math.floor(math.log2(total_num_heads))
|
||||
base = torch.tensor(
|
||||
2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32)
|
||||
slopes = torch.pow(base, powers)
|
||||
|
||||
if closest_power_of_2 != total_num_heads:
|
||||
extra_base = torch.tensor(
|
||||
2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))),
|
||||
dtype=torch.float32,
|
||||
)
|
||||
num_remaining_heads = min(
|
||||
closest_power_of_2, total_num_heads - closest_power_of_2
|
||||
)
|
||||
extra_powers = torch.arange(
|
||||
start=1, end=1 + 2 * num_remaining_heads, step=2, dtype=torch.int32
|
||||
)
|
||||
slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
|
||||
return slopes
|
||||
|
||||
|
||||
def get_alibi_mask(num_heads, seqlen, device, dtype):
|
||||
x = torch.arange(0, seqlen, device=device, dtype=torch.float32).view(-1, 1)
|
||||
y = torch.tensor(seqlen - 1, device=device, dtype=torch.float32).view(1, -1)
|
||||
offsets = -(y - x).view(1, 1, seqlen)
|
||||
return offsets
|
||||
|
||||
|
||||
def ref_masked_attention(
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
scale: float,
|
||||
attn_mask: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
query = query * scale
|
||||
dtype = query.dtype
|
||||
device = query.device
|
||||
query = query.to(torch.float32).cpu()
|
||||
key = key.to(torch.float32).cpu()
|
||||
value = value.to(torch.float32).cpu()
|
||||
attn = torch.einsum("qhd,khd->hqk", query, key)
|
||||
if attn_mask is not None:
|
||||
attn_mask = attn_mask.cpu()
|
||||
attn = attn + attn_mask
|
||||
attn = torch.softmax(attn, dim=-1)
|
||||
out = torch.einsum("hqk,khd->qhd", attn, value)
|
||||
out = out.to(device).to(dtype)
|
||||
return out
|
||||
|
||||
|
||||
def ref_paged_attention_v1(
|
||||
output: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
num_q_per_kv: int,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
use_alibi: bool,
|
||||
) -> None:
|
||||
num_query_heads = query.shape[1]
|
||||
num_kv_heads = value_cache.shape[1]
|
||||
head_size = value_cache.shape[2]
|
||||
block_size = value_cache.shape[3]
|
||||
|
||||
num_input_tokens = query.shape[0]
|
||||
device = output.device
|
||||
slopes = (
|
||||
get_alibi_slopes(num_query_heads)
|
||||
.to(device)
|
||||
.to(torch.float32)
|
||||
.view(num_query_heads, 1, 1)
|
||||
)
|
||||
for i in range(num_input_tokens):
|
||||
q = query[i].unsqueeze(0)
|
||||
block_table = block_tables[i]
|
||||
context_len = int(context_lens[i])
|
||||
|
||||
keys = []
|
||||
values = []
|
||||
for j in range(context_len):
|
||||
block_number = int(block_table[j // block_size])
|
||||
block_offset = j % block_size
|
||||
|
||||
k = key_cache[block_number, :, :, block_offset, :]
|
||||
k = k.reshape(num_kv_heads, head_size)
|
||||
keys.append(k)
|
||||
|
||||
v = value_cache[block_number, :, :, block_offset]
|
||||
values.append(v)
|
||||
keys = torch.stack(keys, dim=0)
|
||||
values = torch.stack(values, dim=0)
|
||||
if num_q_per_kv > 1:
|
||||
keys = torch.repeat_interleave(keys, num_q_per_kv, dim=1)
|
||||
values = torch.repeat_interleave(values, num_q_per_kv, dim=1)
|
||||
scale = 1.0 / (head_size**0.5)
|
||||
if use_alibi:
|
||||
offsets = get_alibi_mask(
|
||||
num_query_heads, context_len, output.device, output.dtype
|
||||
)
|
||||
mask = offsets * slopes
|
||||
mask = mask.to(output.dtype)
|
||||
else:
|
||||
mask = None
|
||||
out = ref_masked_attention(
|
||||
q,
|
||||
keys,
|
||||
values,
|
||||
scale,
|
||||
mask,
|
||||
)
|
||||
out = out.view(num_query_heads, head_size)
|
||||
output[i].copy_(out, non_blocking=True)
|
||||
|
||||
|
||||
def ref_paged_attention_v3(
|
||||
output: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
num_q_per_kv: int,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
use_alibi: bool,
|
||||
) -> None:
|
||||
num_query_heads = query.shape[1]
|
||||
num_kv_heads = value_cache.shape[1]
|
||||
head_size = query.shape[2]
|
||||
block_size = value_cache.shape[2] * 4
|
||||
|
||||
num_input_tokens = query.shape[0]
|
||||
device = output.device
|
||||
slopes = (
|
||||
get_alibi_slopes(num_query_heads)
|
||||
.to(device)
|
||||
.to(torch.float32)
|
||||
.view(num_query_heads, 1, 1)
|
||||
)
|
||||
for i in range(num_input_tokens):
|
||||
q = query[i].unsqueeze(0)
|
||||
block_table = block_tables[i]
|
||||
context_len = int(context_lens[i])
|
||||
|
||||
keys = []
|
||||
values = []
|
||||
for j in range(context_len):
|
||||
block_number = int(block_table[j // block_size])
|
||||
block_offset = j % block_size
|
||||
|
||||
k = key_cache[block_number, :, block_offset // 4, :, block_offset % 4, :]
|
||||
k = k.reshape(num_kv_heads, head_size)
|
||||
keys.append(k)
|
||||
|
||||
v = value_cache[block_number, :, block_offset // 4, :, block_offset % 4, :]
|
||||
v = v.reshape(num_kv_heads, head_size)
|
||||
values.append(v)
|
||||
keys = torch.stack(keys, dim=0)
|
||||
values = torch.stack(values, dim=0)
|
||||
if num_q_per_kv > 1:
|
||||
keys = torch.repeat_interleave(keys, num_q_per_kv, dim=1)
|
||||
values = torch.repeat_interleave(values, num_q_per_kv, dim=1)
|
||||
scale = 1.0 / (head_size**0.5)
|
||||
if use_alibi:
|
||||
offsets = get_alibi_mask(
|
||||
num_query_heads, context_len, output.device, output.dtype
|
||||
)
|
||||
mask = offsets * slopes
|
||||
mask = mask.to(output.dtype)
|
||||
else:
|
||||
mask = None
|
||||
out = ref_masked_attention(
|
||||
q,
|
||||
keys,
|
||||
values,
|
||||
scale,
|
||||
mask,
|
||||
)
|
||||
out = out.view(num_query_heads, head_size)
|
||||
output[i].copy_(out, non_blocking=True)
|
||||
|
||||
|
||||
def rotate_half(x, interleaved=False):
|
||||
if not interleaved:
|
||||
x1, x2 = x.chunk(2, dim=-1)
|
||||
return torch.cat((-x2, x1), dim=-1)
|
||||
else:
|
||||
x1, x2 = x[..., ::2], x[..., 1::2]
|
||||
seq_len, head_nums, _ = x.shape
|
||||
return torch.stack((-x2, x1), dim=-1).reshape(seq_len, head_nums, -1)
|
||||
|
||||
|
||||
def tgi_apply_rotary_emb_torch(
|
||||
x: "torch.Tensor",
|
||||
cos: "torch.Tensor",
|
||||
sin: "torch.Tensor",
|
||||
interleaved: bool = False,
|
||||
):
|
||||
"""
|
||||
x: (seqlen, num_heads, headdim)
|
||||
cos, sin: (seqlen, 1, rotary_dim / 2)
|
||||
interleaved: bool. 在interleaved的实现中,对奇偶维度旋转需要将维度两两交错，实现较为复杂。
|
||||
"""
|
||||
ro_dim = cos.shape[-1] * 2
|
||||
assert ro_dim <= x.shape[-1]
|
||||
assert cos.shape == sin.shape
|
||||
if cos.dim() == 2:
|
||||
cos = cos.unsqueeze(1)
|
||||
sin = sin.unsqueeze(1)
|
||||
if interleaved:
|
||||
cos = cos.repeat_interleave(2, dim=-1)
|
||||
sin = sin.repeat_interleave(2, dim=-1)
|
||||
else:
|
||||
cos = cos.repeat(1, 1, 2)
|
||||
sin = sin.repeat(1, 1, 2)
|
||||
return torch.cat(
|
||||
[
|
||||
x[..., :ro_dim].float() * cos.float()
|
||||
+ rotate_half(x[..., :ro_dim].float(), interleaved) * sin.float(),
|
||||
x[..., ro_dim:].float(),
|
||||
],
|
||||
dim=-1,
|
||||
).to(x.dtype)
|
||||
|
||||
|
||||
def tgi_apply_rotary(
|
||||
querys: List[torch.Tensor],
|
||||
cos: "torch.Tensor",
|
||||
sin: "torch.Tensor",
|
||||
outs: List[torch.Tensor] = None,
|
||||
is_neox_style: bool = True,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
querys: [(num_tokens, num_heads, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16]
|
||||
cos: (max_position, 1, head_size //2) torch.half, torch.float, torch.bfloat16
|
||||
sin: (max_position, 1, head_size //2) torch.half, torch.float, torch.bfloat16
|
||||
is_neox_style: bool
|
||||
判断是否使用Neox,默认为True,即不使用interleaved
|
||||
outs: [(num_tokens, num_heads, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16]
|
||||
Returns:
|
||||
outs: [(num_tokens, num_heads, head_size)] List[torch.half],List[torch.float],List[torch.bfloat16]
|
||||
"""
|
||||
|
||||
assert sin.shape == cos.shape
|
||||
rotary_dim = cos.shape[-1]
|
||||
return_type = False
|
||||
if len(querys) == 1:
|
||||
query = querys[0]
|
||||
query_dim = query.shape[-1]
|
||||
query1 = query[..., :rotary_dim]
|
||||
query2 = query[..., rotary_dim : 2 * rotary_dim]
|
||||
elif len(querys) == 2:
|
||||
return_type = True
|
||||
query1 = querys[0]
|
||||
query2 = querys[1]
|
||||
assert query1.shape == query2.shape
|
||||
query_dim = query1.shape[-1] * 2
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid number for querys: {len(querys)}. " "Expected number 1, or 2."
|
||||
)
|
||||
|
||||
assert rotary_dim * 2 <= query_dim
|
||||
if outs is None:
|
||||
query_shape = query1.shape
|
||||
out = torch.empty(*(query_shape[:-1] + [rotary_dim * 2]))
|
||||
out1 = out[..., :rotary_dim]
|
||||
out2 = out[..., rotary_dim : 2 * rotary_dim]
|
||||
else:
|
||||
assert len(querys) == len(outs)
|
||||
for query, out in zip(querys, outs):
|
||||
assert query.shape == out.shape
|
||||
if len(outs) == 1:
|
||||
out = outs[0]
|
||||
out1 = out[..., :rotary_dim]
|
||||
out2 = out[..., rotary_dim : 2 * rotary_dim]
|
||||
else:
|
||||
out1 = outs[0]
|
||||
out2 = outs[1]
|
||||
|
||||
if cos.dim() == 2:
|
||||
cos = cos.unsqueeze(1)
|
||||
sin = sin.unsqueeze(1)
|
||||
|
||||
ops.infer.tgi_rotary_embedding_neox(
|
||||
query1, query2, cos, sin, out1, out2, is_neox_style
|
||||
)
|
||||
|
||||
if return_type:
|
||||
return out1, out2
|
||||
else:
|
||||
return torch.cat([out1, out2], dim=-1)
|
||||
|
||||
|
||||
def tgi_gather_prefill_logprobs(
|
||||
logits: "torch.Tensor",
|
||||
prefill_tokens_indices: "torch.Tensor",
|
||||
output: "torch.Tensor" = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
logits: (num_tokens, vocab_size) torch.half, torch.bfloat16
|
||||
prefill_tokens_indices: (tokens_indices) torch.int
|
||||
output: (tokens_indices, 1) torch.half, torch.bfloat16
|
||||
Returns:
|
||||
output: (tokens_indices, 1) torch.half, torch.bfloat16
|
||||
"""
|
||||
if output is None:
|
||||
output = logits.new_empty(prefill_tokens_indices.shape)
|
||||
ops.infer.tgi_gather_prefill_logprobs(logits, prefill_tokens_indices, output)
|
||||
return output
|
||||
|
||||
|
||||
def paged_attention_v1(
|
||||
output: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
num_kv_heads: torch.Tensor,
|
||||
scale: float,
|
||||
block_tables: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
block_size: int,
|
||||
max_context_len: int,
|
||||
alibi_slopes: torch.Tensor = None,
|
||||
use_sqrt_alibi: bool = False,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
query: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
key_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16
|
||||
value_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16
|
||||
num_kv_heads: int
|
||||
scale: float
|
||||
block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64
|
||||
context_lens_cpu: (num_tokens) torch.int32
|
||||
context_lens: (num_tokens) torch.int32
|
||||
block_size: int
|
||||
max_context_len: int
|
||||
alibi_slopes: (num_heads) torch.float32
|
||||
use_sqrt_alibi: bool
|
||||
Returns:
|
||||
output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
"""
|
||||
ops.infer.tgi_single_query_cached_kv_attention(
|
||||
output,
|
||||
query,
|
||||
key_cache,
|
||||
value_cache,
|
||||
num_kv_heads,
|
||||
scale,
|
||||
block_tables,
|
||||
context_lens,
|
||||
block_size,
|
||||
max_context_len,
|
||||
query.stride(0),
|
||||
use_sqrt_alibi,
|
||||
alibi_slopes,
|
||||
)
|
||||
|
||||
|
||||
def paged_attention_v3(
|
||||
output: "torch.Tensor",
|
||||
query: "torch.Tensor",
|
||||
key_cache: "torch.Tensor",
|
||||
value_cache: "torch.Tensor",
|
||||
head_mapping: "torch.Tensor",
|
||||
scale: float,
|
||||
block_tables: "torch.Tensor",
|
||||
context_lens: "torch.Tensor",
|
||||
block_size: int,
|
||||
max_context_len: int,
|
||||
alibi_slopes: "torch.Tensor" = None,
|
||||
use_sqrt_alibi: bool = False,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
query: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
key_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16
|
||||
value_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16
|
||||
num_kv_heads: int
|
||||
scale: float
|
||||
block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64
|
||||
context_lens_cpu: (num_tokens) torch.int32
|
||||
context_lens: (num_tokens) torch.int32
|
||||
block_size: int
|
||||
max_context_len: int
|
||||
alibi_slopes: (num_heads) torch.float32
|
||||
use_sqrt_alibi: bool
|
||||
Returns:
|
||||
output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
"""
|
||||
ops.infer.single_query_cached_kv_attention_v3(
|
||||
output,
|
||||
query,
|
||||
key_cache,
|
||||
value_cache,
|
||||
head_mapping,
|
||||
scale,
|
||||
block_tables,
|
||||
context_lens,
|
||||
block_size,
|
||||
max_context_len,
|
||||
query.stride(0),
|
||||
use_sqrt_alibi,
|
||||
alibi_slopes,
|
||||
)
|
||||
|
||||
|
||||
|
||||
def paged_attention_v7(
|
||||
output: torch.Tensor,
|
||||
query: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
num_kv_heads: int,
|
||||
scale: float,
|
||||
block_tables: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
block_size: int,
|
||||
max_context_len: int,
|
||||
alibi_slopes: torch.Tensor = None,
|
||||
use_sqrt_alibi: bool = False,
|
||||
):
|
||||
|
||||
"""
|
||||
Args:
|
||||
output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
query: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
key_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16
|
||||
value_cache: (num_blocks, num_kv_heads, block_size, head_size) torch.half, torch.bfloat16
|
||||
num_kv_heads: int
|
||||
scale: float
|
||||
block_tables: (num_tokens, max_num_blocks_per_seq) torch.int64
|
||||
context_lens: (num_tokens) torch.int32
|
||||
block_size: int
|
||||
max_context_len: int
|
||||
alibi_slopes: (num_heads) torch.float32
|
||||
use_sqrt_alibi: bool
|
||||
Returns:
|
||||
output: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
"""
|
||||
num_blocks = key_cache.size(0)
|
||||
head_size = query.size(-1)
|
||||
key_cache = key_cache.view(num_blocks, num_kv_heads, block_size, head_size)
|
||||
value_cache = value_cache.view(num_blocks, num_kv_heads, block_size, head_size)
|
||||
ops.infer.vllm_paged_attention(
|
||||
output,
|
||||
query,
|
||||
key_cache,
|
||||
value_cache,
|
||||
num_kv_heads,
|
||||
scale,
|
||||
block_tables,
|
||||
context_lens,
|
||||
block_size,
|
||||
max_context_len,
|
||||
alibi_slopes,
|
||||
True,
|
||||
-1,
|
||||
-1,
|
||||
0.0,
|
||||
False,
|
||||
use_sqrt_alibi,
|
||||
)
|
||||
return output
|
||||
|
||||
def reshape_and_cache_v1(
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
slot_mapping: torch.Tensor,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
key: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16
|
||||
value: (num_tokens, num_heads, head_size) torch.half, torch.float, torch.bfloat16
|
||||
key_cache: (num_blocks, num_heads, head_size//8, block_size, 8) or (num_blocks, num_heads, head_size//4, block_size, 4) torch.half, torch.float, torch.bfloat16
|
||||
if dtype=torch.half or torch.bfloat16,key_cache shape: (num_blocks, num_heads, head_size//8, block_size, 8)
|
||||
if dtype=torch.float,key_cache shape: (num_blocks, num_heads, head_size//4, block_size, 4)
|
||||
value_cache: (num_blocks, num_heads, head_size//8, block_size, 8) or (num_blocks, num_heads, head_size//4, block_size, 4) torch.half, torch.float, torch.bfloat16
|
||||
if dtype=torch.half or torch.bfloat16,value_cache shape: (num_blocks, num_heads, head_size//8, block_size, 8)
|
||||
if dtype=torch.float,value_cache shape: (num_blocks, num_heads, head_size//4, block_size, 4)
|
||||
slot_mapping: (num_tokens) torch.long
|
||||
Returns:
|
||||
None, 对key_cache,value_cache进行in place 操作
|
||||
"""
|
||||
ops.infer.vllm_cache_ops_reshape_and_cache_v4(
|
||||
key,
|
||||
value,
|
||||
key_cache,
|
||||
value_cache,
|
||||
slot_mapping,
|
||||
key.stride(0),
|
||||
value.stride(0),
|
||||
)
|
||||
|
||||
|
||||
def reshape_and_cache(
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
key_cache: torch.Tensor,
|
||||
value_cache: torch.Tensor,
|
||||
slot_mapping: torch.Tensor,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
key: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
value: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
key_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.bfloat16
|
||||
value_cache: (num_blocks, num_heads, block_size, head_size) torch.half, torch.bfloat16
|
||||
slot_mapping: (num_tokens) torch.long
|
||||
Returns:
|
||||
None, 对key_cache,value_cache进行in place 操作
|
||||
"""
|
||||
num_tokens, num_kv_heads, head_size = key.shape
|
||||
num_blocks = key_cache.size(0)
|
||||
key_cache = key_cache.view(num_blocks, num_kv_heads, -1, head_size)
|
||||
value_cache = value_cache.view(num_blocks, num_kv_heads, -1, head_size)
|
||||
ops.infer.vllm_cache_ops_reshape_and_cache(
|
||||
key,
|
||||
value,
|
||||
key_cache,
|
||||
value_cache,
|
||||
slot_mapping,
|
||||
key.stride(0),
|
||||
value.stride(0),
|
||||
)
|
||||
|
||||
|
||||
def ref_reshape_and_cache_v3(
|
||||
key,
|
||||
value,
|
||||
key_cache,
|
||||
value_cache,
|
||||
slot_mapping,
|
||||
num_tokens,
|
||||
num_heads,
|
||||
head_size,
|
||||
block_size,
|
||||
):
|
||||
reshaped_key = key.view(num_tokens, num_heads, head_size // 32, 32)
|
||||
reshaped_value = value.reshape(num_tokens, num_heads, head_size // 32, 32)
|
||||
for i in range(num_tokens):
|
||||
|
||||
block_idx = torch.div(slot_mapping[i], block_size, rounding_mode="floor")
|
||||
block_offset = slot_mapping[i] % block_size
|
||||
|
||||
key_cache[
|
||||
block_idx, :, block_offset // 4, :, block_offset % 4, :
|
||||
] = reshaped_key[i]
|
||||
value_cache[
|
||||
block_idx, :, block_offset // 4, :, block_offset % 4, :
|
||||
] = reshaped_value[i]
|
||||
|
||||
|
||||
def reshape_and_cache_v3(
|
||||
key: "torch.Tensor",
|
||||
value: "torch.Tensor",
|
||||
key_cache: "torch.Tensor",
|
||||
value_cache: "torch.Tensor",
|
||||
slot_mapping: "torch.Tensor",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
key: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
value: (num_tokens, num_heads, head_size) torch.half, torch.bfloat16
|
||||
key_cache: (num_blocks, num_heads, block_size // 4, head_size // 32, 4, 32) torch.half, torch.bfloat16
|
||||
目前block_size 只支持16,head_size 只支持64,128,256
|
||||
value_cache: (num_blocks, num_heads, block_size // 4, head_size // 32, 4, 32) torch.half, torch.bfloat16
|
||||
slot_mapping: (num_tokens) torch.int
|
||||
Returns:
|
||||
None, 对key_cache,value_cache进行in place 操作
|
||||
"""
|
||||
if key.dim() != 3 or key.shape != value.shape or key.size(-1) not in [64, 128, 256]:
|
||||
raise NotImplementedError(
|
||||
"reshape_and_cache_v3 only support key.dim()==3 and key.shape== value.shape and head_size must be 64, 128 , 256!"
|
||||
)
|
||||
ops.infer.cache_ops_reshape_and_cache_v3(
|
||||
key,
|
||||
value,
|
||||
key_cache,
|
||||
value_cache,
|
||||
slot_mapping,
|
||||
key.stride(0),
|
||||
value.stride(0),
|
||||
)
|
||||
2033
ixformer_sdk/inference/functions/vllm.py
Normal file
2033
ixformer_sdk/inference/functions/vllm.py
Normal file
File diff suppressed because it is too large
Load Diff
225
ixformer_sdk/inference/functions/w8a16.py
Normal file
225
ixformer_sdk/inference/functions/w8a16.py
Normal file
@@ -0,0 +1,225 @@
|
||||
import math
|
||||
from typing import List, Union, Optional
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
from torch.autograd.function import Function, FunctionCtx
|
||||
|
||||
import ixformer
|
||||
from ixformer.core import config
|
||||
|
||||
__all__ = [
|
||||
"w8a16_gemm",
|
||||
"w8a16_gemv",
|
||||
"w8a16",
|
||||
"ref_w8a16",
|
||||
"wu8a16",
|
||||
"ref_wu8a16",
|
||||
]
|
||||
|
||||
|
||||
def w8a16_gemv(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
group_size: int = -1,
|
||||
format: str = "unknown",
|
||||
output: Optional[torch.Tensor] = None
|
||||
):
|
||||
"""
|
||||
w8a16 gemv 接口
|
||||
input : bf16|fp16 (bs, ic)
|
||||
qweights : int8 TN:(oc, ic) NN:(ic, oc)
|
||||
scales : bf16|fp16 TN: 当groupsize为-1时, shape: (1, oc), 否则,shape: (ic/group_size, oc) NN:(1, oc)
|
||||
TN 支持条件： ic % groupSize = 0, oc % 2 = 0, bs<=4
|
||||
NN 支持条件： groupsize = -1 or groupsize = ic, oc % 4 = 0, bs<=4
|
||||
"""
|
||||
assert format in ["TN", "NN"]
|
||||
assert len(qweights.shape) == 2
|
||||
assert len(scales.shape) == 2
|
||||
|
||||
input_shape = list(inputs.shape)
|
||||
inputs = inputs.view(-1, input_shape[-1])
|
||||
|
||||
if format == "TN":
|
||||
output_shape = input_shape[:-1] + [qweights.shape[0]]
|
||||
else:
|
||||
output_shape = input_shape[:-1] + [qweights.shape[1]]
|
||||
|
||||
if output is None:
|
||||
output = inputs.new_empty(output_shape).view(-1, output_shape[-1])
|
||||
|
||||
ops.infer.w8a16_gemv(output, inputs, qweights, scales, group_size, format)
|
||||
return output.view(output_shape)
|
||||
|
||||
|
||||
def w8a16_gemm(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
group_size: int = -1,
|
||||
format: str = "TN",
|
||||
persistent: int = 0,
|
||||
output: Optional[torch.Tensor] = None
|
||||
):
|
||||
"""
|
||||
w8a16 gemm 接口
|
||||
1. group_size=-1 or group_size=ic
|
||||
input : bf16|fp16 (bs, ic)
|
||||
qweights : int8 TN:(oc, ic) NN:(ic, oc)
|
||||
scales : bf16|fp16 (1, oc)
|
||||
NN 支持条件： ic%64==0, oc%64==0
|
||||
|
||||
2. group_size=64
|
||||
input : bf16|fp16 (bs, ic)
|
||||
qweights : int8 TN:(oc, ic)
|
||||
scales : bf16|fp16 (ic/64, oc)
|
||||
TN 支持条件： oc%2==0, ic%64==0
|
||||
NN 不支持
|
||||
"""
|
||||
|
||||
assert format in ["TN", "NN"]
|
||||
assert len(qweights.shape) == 2
|
||||
assert len(scales.shape) == 2
|
||||
|
||||
input_shape = list(inputs.shape)
|
||||
inputs = inputs.view(-1, input_shape[-1])
|
||||
|
||||
if format == "TN":
|
||||
output_shape = input_shape[:-1] + [qweights.shape[0]]
|
||||
else:
|
||||
output_shape = input_shape[:-1] + [qweights.shape[1]]
|
||||
|
||||
if output is None:
|
||||
output = inputs.new_empty(output_shape).view(-1, output_shape[-1])
|
||||
|
||||
ops.infer.w8a16_gemm(
|
||||
output, inputs, qweights, scales, group_size, format, persistent
|
||||
)
|
||||
return output.view(output_shape)
|
||||
|
||||
|
||||
def dequant(qweight, scales, group_size):
|
||||
IC, OC = qweight.shape
|
||||
weight = qweight.t().reshape(OC, -1, group_size).to(
|
||||
torch.float32
|
||||
) * scales.t().unsqueeze(-1)
|
||||
return weight.reshape(OC, IC)
|
||||
|
||||
|
||||
def ref_w8a16(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
group_size: int = -1,
|
||||
format: str = "TN",
|
||||
):
|
||||
if group_size == -1:
|
||||
group_size = inputs.shape[1]
|
||||
if format == "TN":
|
||||
weights = dequant(qweights.transpose(0, 1), scales, group_size)
|
||||
elif format == "NN":
|
||||
weights = dequant(qweights, scales, group_size)
|
||||
return torch.nn.functional.linear(inputs, weights.to(inputs.dtype))
|
||||
|
||||
|
||||
def w8a16(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
group_size: int = -1,
|
||||
format: str = "TN",
|
||||
output: Optional[torch.Tensor] = None,
|
||||
persistent: int = 0,
|
||||
):
|
||||
input_shape = inputs.shape
|
||||
inputs = inputs.view(-1, input_shape[-1])
|
||||
bs = inputs.size(0)
|
||||
inputs = inputs.view(input_shape)
|
||||
if bs <= config.IXFORMER_GEMV_THRESHOLD:
|
||||
return w8a16_gemv(
|
||||
inputs=inputs,
|
||||
qweights=qweights,
|
||||
scales=scales,
|
||||
group_size=group_size,
|
||||
format=format,
|
||||
output=output
|
||||
)
|
||||
else:
|
||||
return w8a16_gemm(
|
||||
inputs=inputs,
|
||||
qweights=qweights,
|
||||
scales=scales,
|
||||
group_size=group_size,
|
||||
format=format,
|
||||
output=output,
|
||||
persistent=persistent
|
||||
)
|
||||
|
||||
|
||||
def ref_wu8a16(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
zeros: "torch.Tensor",
|
||||
group_size: int = -1,
|
||||
format: str = "TN",
|
||||
):
|
||||
assert format in ["TN"]
|
||||
assert len(qweights.shape) == 2
|
||||
assert len(scales.shape) == 2
|
||||
org_w_shape = qweights.shape
|
||||
scales = scales.transpose(0, 1).flatten().view(-1, 1)
|
||||
zeros = zeros.transpose(0, 1).flatten().view(-1, 1)
|
||||
|
||||
if group_size != -1:
|
||||
qweights = qweights.reshape(-1, group_size)
|
||||
w = (qweights - zeros) * scales
|
||||
w = w.reshape(org_w_shape)
|
||||
output = torch.matmul(inputs, w.t())
|
||||
return output
|
||||
|
||||
|
||||
def wu8a16(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
zeros: "torch.Tensor",
|
||||
group_size: int = -1,
|
||||
format: str = "TN",
|
||||
persistent: int = 0,
|
||||
):
|
||||
"""
|
||||
http://confluence.iluvatar.ai:8090/display/SW/cuinferCustomGemm+Interface+Doc
|
||||
wu8a16 非对称量化 gemm 接口
|
||||
1. group_size=-1
|
||||
input : bf16|fp16 (bs, ic)
|
||||
qweights : uint8 TN:(oc, ic)
|
||||
scales : bf16|fp16 (1,oc)
|
||||
zeros : bf16|fp16 (1,oc)
|
||||
TN 支持条件： ic % 64 == 0
|
||||
NN 不支持
|
||||
|
||||
2. group_size=64
|
||||
input : bf16|fp16 (bs, ic)
|
||||
qweights : int8 TN:(oc, ic)
|
||||
scales : bf16|fp16 (ic/64, oc)
|
||||
zeros : bf16|fp16 (ic/64, oc)
|
||||
TN 支持条件： oc % 2 == 0 && ic % 64 == 0
|
||||
NN 不支持
|
||||
"""
|
||||
assert format in ["TN"]
|
||||
assert len(qweights.shape) == 2
|
||||
assert len(scales.shape) == 2
|
||||
|
||||
input_shape = list(inputs.shape)
|
||||
inputs = inputs.view(-1, input_shape[-1])
|
||||
|
||||
output_shape = input_shape[:-1] + [qweights.shape[0]]
|
||||
|
||||
output = inputs.new_empty(output_shape).view(-1, output_shape[-1])
|
||||
|
||||
ops.infer.wu8a16_gemm(
|
||||
output, inputs, qweights, scales, zeros, group_size, format, persistent
|
||||
)
|
||||
return output.view(output_shape)
|
||||
317
ixformer_sdk/inference/functions/w8a8.py
Normal file
317
ixformer_sdk/inference/functions/w8a8.py
Normal file
@@ -0,0 +1,317 @@
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = [
|
||||
"w8a8",
|
||||
"ref_w8a8",
|
||||
"dynamic_scaled_int8_quant",
|
||||
"ref_dynamic_scaled_int8_quant",
|
||||
"static_scaled_int8_quant",
|
||||
"ref_static_scaled_int8_quant",
|
||||
"scaled_int8_quant",
|
||||
]
|
||||
|
||||
|
||||
def ref_w8a8(
|
||||
input: "torch.Tensor",
|
||||
weight: "torch.Tensor",
|
||||
i_scales: "torch.Tensor",
|
||||
w_scales: "torch.Tensor",
|
||||
output: "torch.Tensor",
|
||||
format: str = "TN",
|
||||
persistent=0,
|
||||
bias: torch.Tensor = None,
|
||||
):
|
||||
dtype = output.dtype
|
||||
input_f32 = input.to(torch.float32)
|
||||
weight_f32 = weight.to(torch.float32)
|
||||
assert format in ["TN", "NN", "NT"]
|
||||
if format == "TN":
|
||||
weight_f32 = weight_f32.transpose(0, 1)
|
||||
if format == "NT":
|
||||
input_f32 = input_f32.transpose(0, 1)
|
||||
output_f32 = (
|
||||
torch.matmul(input_f32, weight_f32)
|
||||
* i_scales.view(-1, 1)
|
||||
* w_scales.view(1, -1)
|
||||
)
|
||||
if bias is not None:
|
||||
bias_f32 = bias.to(torch.float32)
|
||||
output_f32 += bias_f32.view(1, -1)
|
||||
output.copy_(output_f32.to(dtype))
|
||||
return output
|
||||
|
||||
|
||||
def w8a8_gemm(
|
||||
input: "torch.Tensor",
|
||||
weight: "torch.Tensor",
|
||||
i_scales: "torch.Tensor",
|
||||
w_scales: "torch.Tensor",
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
output: Optional[torch.Tensor] = None,
|
||||
format: str = "TN",
|
||||
persistent: bool = False,
|
||||
out_dtype: torch.dtype = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (n, k) torch.int8
|
||||
weight: (m, k) if format == "TN" else (k, m) torch.int8
|
||||
i_scales: (n) torch.float32
|
||||
w_scales: (m) torch.float32
|
||||
bias: (m) torch.float32, same as output_type
|
||||
format: str
|
||||
Options include TN, NN and NT
|
||||
persistent: Whether to use overleap bool
|
||||
out_dtype: torch.float16, torch.bfloat16
|
||||
Returns:
|
||||
output: (n, m) torch.float16, torch.bfloat16
|
||||
"""
|
||||
|
||||
input_shape = input.shape
|
||||
|
||||
if output is None:
|
||||
if out_dtype is None:
|
||||
raise RuntimeError("w8a8 gemm need out_dtype argument when output is none.")
|
||||
output = torch.empty(
|
||||
(input_shape[:-1] + (weight.shape[0],)),
|
||||
dtype=out_dtype,
|
||||
device=input.device,
|
||||
)
|
||||
|
||||
output_shape = output.shape
|
||||
|
||||
input = input.view(-1, input_shape[-1])
|
||||
output = output.view(-1, output_shape[-1])
|
||||
|
||||
ops.infer.w8a8_gemm(
|
||||
output, input, weight, i_scales, w_scales, bias, format, int(persistent)
|
||||
)
|
||||
|
||||
return output.view(*output_shape)
|
||||
|
||||
|
||||
def ref_static_scaled_int8_quant(output, input, scale):
|
||||
"""
|
||||
Args:
|
||||
output: [torch.int8] [m, k]
|
||||
input: [torch.half,torch.bfloat16] [m, k]
|
||||
scale: [torch.float32] [1]
|
||||
Returns:
|
||||
output: [torch.int8] [m, k]
|
||||
scale: [torch.float32] [1]
|
||||
"""
|
||||
# [m, 1]
|
||||
f_input = input / scale.to(input.dtype)
|
||||
i_output = torch.clamp(torch.round(f_input), -127, 127).to(torch.int8)
|
||||
output.copy_(i_output)
|
||||
return output, scale
|
||||
|
||||
|
||||
# for vllm: https://github.com/vllm-project/vllm/blob/v0.5.4/vllm/_custom_ops.py#L387
|
||||
def static_scaled_int8_quant(output, input, scale):
|
||||
"""
|
||||
Args:
|
||||
output: [torch.int8] [m, k]
|
||||
input: [torch.half,torch.bfloat16] [m, k]
|
||||
scale: [torch.float32] [1]
|
||||
Returns:
|
||||
output: [torch.int8] [m, k]
|
||||
scale: [torch.float32] [1]
|
||||
"""
|
||||
ops.infer.scaled_int8_quant(output, input, scale, 0)
|
||||
return output, scale
|
||||
|
||||
|
||||
def ref_dynamic_scaled_int8_quant(output, input, scale):
|
||||
"""
|
||||
Args:
|
||||
output: [torch.int8] [m, k]
|
||||
input: [torch.half,torch.bfloat16] [m, k]
|
||||
scale: [torch.float32] [m]
|
||||
Returns:
|
||||
output: [torch.int8] [m, k]
|
||||
scale: [torch.float32] [m]
|
||||
"""
|
||||
# [m, 1]
|
||||
amax_, _ = torch.max(torch.abs(input), dim=-1, keepdim=True)
|
||||
f_scale = amax_.float() / 127.0
|
||||
scale.view(-1).copy_(f_scale.view(-1))
|
||||
|
||||
f_input = input / f_scale.to(input.dtype)
|
||||
i_output = torch.clamp(torch.round(f_input), -127, 127).to(torch.int8)
|
||||
output.copy_(i_output)
|
||||
return output, scale.view(input.shape[:-1])
|
||||
|
||||
|
||||
# for vllm: https://github.com/vllm-project/vllm/blob/v0.5.4/vllm/_custom_ops.py#L394
|
||||
def dynamic_scaled_int8_quant(output, input, scale):
|
||||
"""
|
||||
Args:
|
||||
output: [torch.int8] [m, k]
|
||||
input: [torch.half,torch.bfloat16] [m, k]
|
||||
scale: [torch.float32] [m]
|
||||
Returns:
|
||||
output: [torch.int8] [m, k]
|
||||
scale: [torch.float32] [m]
|
||||
"""
|
||||
ops.infer.scaled_int8_quant(output, input, scale, 1)
|
||||
return output, scale
|
||||
|
||||
|
||||
def scaled_int8_quant(
|
||||
input: torch.Tensor, scale: Optional[torch.Tensor] = None
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Quantize the input tensor to int8 and return the quantized tensor and scale.
|
||||
|
||||
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.
|
||||
|
||||
Returns:
|
||||
Tuple[Torch.Tensor, Torch.Tensor] : Output int8 tensor and scales.
|
||||
"""
|
||||
output = torch.empty_like(input, dtype=torch.int8)
|
||||
if scale is not None:
|
||||
# static-per-tensor quantization.
|
||||
static_scaled_int8_quant(output, input, scale)
|
||||
return output, scale
|
||||
|
||||
# dynamic-per-token quantization.
|
||||
input_scales = torch.empty(
|
||||
(input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32
|
||||
)
|
||||
dynamic_scaled_int8_quant(output, input, input_scales)
|
||||
return output, input_scales
|
||||
|
||||
|
||||
def w8a8_gemv(
|
||||
input: "torch.Tensor",
|
||||
weight: "torch.Tensor",
|
||||
i_scales: "torch.Tensor",
|
||||
w_scales: "torch.Tensor",
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
output: Optional[torch.Tensor] = None,
|
||||
format: str = "TN",
|
||||
persistent: bool = False,
|
||||
out_dtype: torch.dtype = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (n, k) torch.int8
|
||||
weight: (m, k) if format == "TN" else (k, m) torch.int8
|
||||
i_scales: (n) torch.float32
|
||||
w_scales: (m) torch.float32
|
||||
bias: (m) torch.float32 same as output_type
|
||||
format: str
|
||||
Options include TN and NN
|
||||
persistent: Whether to use overleap bool
|
||||
out_dtype: torch.float16, torch.bfloat16
|
||||
Returns:
|
||||
output: (n, m) torch.float16, torch.bfloat16
|
||||
"""
|
||||
|
||||
input_shape = input.shape
|
||||
|
||||
if output is None:
|
||||
if out_dtype is None:
|
||||
raise RuntimeError("w8a8 gemv need out_dtype argument when output is none.")
|
||||
output = torch.empty(
|
||||
(input_shape[:-1] + (weight.shape[0],)),
|
||||
dtype=out_dtype,
|
||||
device=input.device,
|
||||
)
|
||||
|
||||
input = input.view(-1, input_shape[-1])
|
||||
|
||||
ops.infer.w8a8_gemv(
|
||||
output, input, weight, i_scales, w_scales, bias, format, int(persistent)
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def handle_pading(weight: torch.Tensor, format: str, is_gemm: bool):
|
||||
"""Handle padding alignment for weight matrices
|
||||
Args:
|
||||
weight: Original weight matrix [m, k]
|
||||
format: Matrix format, TN indicates transposed layout
|
||||
is_gemm: Whether for GEMM operation (requires extra alignment checks)
|
||||
Returns:
|
||||
torch.Tensor: Padded weight matrix
|
||||
Raises:
|
||||
AssertionError: When is_gemm=True requires 4-byte alignment for m/k
|
||||
"""
|
||||
# weight should have been pad before w8a8 is called, handle _padding here just ensure the code run success,
|
||||
# but performance is low, please refer to vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
|
||||
m, k = weight.shape
|
||||
s = weight.stride(0)
|
||||
if s % 64 != 0 and format == "TN":
|
||||
pad_k = (s // 64 + 1) * 64
|
||||
weight_pad = torch.empty((m, pad_k), dtype=weight.dtype, device=weight.device)
|
||||
_weight = weight_pad[:, :k]
|
||||
if is_gemm:
|
||||
assert m % 4 == 0 and k % 4 == 0
|
||||
_weight.copy_(weight)
|
||||
return _weight
|
||||
else:
|
||||
return weight
|
||||
|
||||
|
||||
def w8a8(
|
||||
input: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
i_scales: torch.Tensor,
|
||||
w_scales: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
output: Optional[torch.Tensor] = None,
|
||||
format: str = "TN",
|
||||
persistent: bool = False,
|
||||
out_dtype: torch.dtype = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
input: (n, k) torch.int8
|
||||
weight: (m, k) if format == "TN" else (k, m) torch.int8
|
||||
i_scales: (n) torch.float32
|
||||
w_scales: (m) torch.float32
|
||||
bias: (m) torch.float32, same as output_type
|
||||
format: str
|
||||
Options include TN and NN
|
||||
persistent: Whether to use overleap bool
|
||||
out_dtype: torch.float16, torch.bfloat16
|
||||
Returns:
|
||||
output: (n, m) torch.float16, torch.bfloat16
|
||||
"""
|
||||
bs = input.numel() // input.shape[-1]
|
||||
gemv_condition = (format == "TN" and bs <= 1) or (format == "NN" and bs <= 16)
|
||||
if gemv_condition:
|
||||
weight = handle_pading(weight, format, is_gemm=False)
|
||||
return w8a8_gemv(
|
||||
input,
|
||||
weight,
|
||||
i_scales,
|
||||
w_scales,
|
||||
bias=bias,
|
||||
output=output,
|
||||
format=format,
|
||||
persistent=persistent,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
else:
|
||||
weight = handle_pading(weight, format, is_gemm=True)
|
||||
return w8a8_gemm(
|
||||
input,
|
||||
weight,
|
||||
i_scales,
|
||||
w_scales,
|
||||
bias=bias,
|
||||
output=output,
|
||||
format=format,
|
||||
persistent=persistent,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
157
ixformer_sdk/inference/functions/wi4a16.py
Normal file
157
ixformer_sdk/inference/functions/wi4a16.py
Normal file
@@ -0,0 +1,157 @@
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = ["wi4a16_gemm", "wi4a16_gemv", "wi4a16", "ref_wi4a16"]
|
||||
|
||||
|
||||
def dequant_weight(tensor, scales, zeros, block_size):
|
||||
# from CPM
|
||||
"""
|
||||
tensor: (oc/2, ic)
|
||||
scales: (oc, ic/group_size)
|
||||
zeros: (oc, ic/group_size)
|
||||
"""
|
||||
dtype = scales.dtype
|
||||
left = tensor >> 4
|
||||
right = tensor << 4 >> 4
|
||||
left, right = right, left
|
||||
ret = torch.cat((left, right), dim=-1).reshape(-1, left.size(-1))
|
||||
ret_shape = ret.size()
|
||||
ret = ret.view(-1, block_size)
|
||||
ret = scales.view(-1, 1) * (ret - zeros.view(-1, 1))
|
||||
ret = ret.reshape(ret_shape).to(dtype=dtype)
|
||||
return ret
|
||||
|
||||
|
||||
def ref_wi4a16(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
zeros: "torch.Tensor",
|
||||
group_size: int = -1,
|
||||
format: str = "TN",
|
||||
):
|
||||
assert format in ["TN"]
|
||||
weights = dequant_weight(
|
||||
qweights,
|
||||
scales.transpose(0, 1).contiguous(),
|
||||
zeros.transpose(0, 1).contiguous(),
|
||||
group_size,
|
||||
)
|
||||
output = torch.nn.functional.linear(inputs, weights.to(inputs.dtype))
|
||||
return output
|
||||
|
||||
|
||||
def wi4a16_gemm(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
zeros: "torch.Tensor",
|
||||
group_size: int = -1,
|
||||
format: str = "TN",
|
||||
output=None,
|
||||
):
|
||||
"""
|
||||
wi4a16 gemm 接口
|
||||
支持条件:
|
||||
format = TN
|
||||
group_size = 128
|
||||
input : fp16 (bs, ic)
|
||||
qweights : int8 (oc/2, ic)
|
||||
scales : fp16 (ic/group_size, oc)
|
||||
zeros : fp16 (ic/group_size, oc)
|
||||
TN 支持条件： oc % 2 == 0 && ic % 128 == 0
|
||||
NN 支持条件： 不支持
|
||||
"""
|
||||
|
||||
assert format in ["TN"]
|
||||
assert len(qweights.shape) == 2
|
||||
assert len(scales.shape) == 2
|
||||
assert len(zeros.shape) == 2
|
||||
|
||||
input_shape = list(inputs.shape)
|
||||
inputs = inputs.view(-1, input_shape[-1])
|
||||
|
||||
if output is None:
|
||||
output_shape = input_shape[:-1] + [scales.shape[1]]
|
||||
output = inputs.new_empty(output_shape).view(-1, output_shape[-1])
|
||||
else:
|
||||
output_shape = output.shape
|
||||
|
||||
ops.infer.wi4a16_gemm(output, inputs, qweights, scales, zeros, group_size, format)
|
||||
return output.view(output_shape)
|
||||
|
||||
|
||||
def wi4a16_gemv(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
zeros: "torch.Tensor",
|
||||
group_size: int = -1,
|
||||
format: str = "TN",
|
||||
output=None,
|
||||
):
|
||||
"""
|
||||
wi4a16 gemv 接口
|
||||
支持条件:
|
||||
format = TN
|
||||
group_size = 128
|
||||
input : bf16|fp16 (bs, ic)
|
||||
qweights : int8 (oc/2, ic)
|
||||
scales : bf16|fp16 (ic/group_size, oc)
|
||||
zeros : bf16|fp16 (ic/group_size, oc)
|
||||
TN 支持条件： oc % 2 == 0 && ic % 128 == 0
|
||||
NN 支持条件： 不支持
|
||||
"""
|
||||
|
||||
assert format in ["TN"]
|
||||
assert len(qweights.shape) == 2
|
||||
assert len(scales.shape) == 2
|
||||
assert len(zeros.shape) == 2
|
||||
|
||||
input_shape = list(inputs.shape)
|
||||
inputs = inputs.view(-1, input_shape[-1])
|
||||
|
||||
if output is None:
|
||||
output_shape = input_shape[:-1] + [scales.shape[1]]
|
||||
output = inputs.new_empty(output_shape).view(-1, output_shape[-1])
|
||||
else:
|
||||
output_shape = output.shape
|
||||
|
||||
ops.infer.wi4a16_gemv(output, inputs, qweights, scales, zeros, group_size, format)
|
||||
return output.view(output_shape)
|
||||
|
||||
|
||||
def wi4a16(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
zeros: "torch.Tensor",
|
||||
group_size: int = -1,
|
||||
format: str = "TN",
|
||||
output=None,
|
||||
):
|
||||
input_shape = inputs.shape
|
||||
inputs = inputs.view(-1, input_shape[-1])
|
||||
bs = inputs.size(0)
|
||||
inputs = inputs.view(input_shape)
|
||||
if bs <= 1:
|
||||
return wi4a16_gemv(
|
||||
inputs=inputs,
|
||||
qweights=qweights,
|
||||
scales=scales,
|
||||
zeros=zeros,
|
||||
group_size=group_size,
|
||||
format=format,
|
||||
output=output,
|
||||
)
|
||||
else:
|
||||
return wi4a16_gemm(
|
||||
inputs=inputs,
|
||||
qweights=qweights,
|
||||
scales=scales,
|
||||
zeros=zeros,
|
||||
group_size=group_size,
|
||||
format=format,
|
||||
output=output,
|
||||
)
|
||||
155
ixformer_sdk/inference/functions/wui4a16.py
Normal file
155
ixformer_sdk/inference/functions/wui4a16.py
Normal file
@@ -0,0 +1,155 @@
|
||||
import ixformer._C as ops
|
||||
import torch
|
||||
|
||||
__all__ = ["wui4a16_gemm", "wui4a16_gemv", "wui4a16", "ref_wui4a16"]
|
||||
|
||||
|
||||
def dequant_weight(tensor, scales, zeros, block_size):
|
||||
"""
|
||||
tensor: (oc/2, ic)
|
||||
scales: (oc, ic/group_size)
|
||||
zeros: (oc, ic/group_size)
|
||||
"""
|
||||
dtype = scales.dtype
|
||||
left = tensor >> 4
|
||||
right = tensor << 4 >> 4
|
||||
left, right = right, left
|
||||
ret = torch.cat((left, right), dim=-1).reshape(-1, left.size(-1))
|
||||
ret_shape = ret.size()
|
||||
ret = ret.view(-1, block_size)
|
||||
ret = scales.view(-1, 1) * (ret - zeros.view(-1, 1))
|
||||
ret = ret.reshape(ret_shape).to(dtype=dtype)
|
||||
return ret
|
||||
|
||||
|
||||
def ref_wui4a16(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
zeros: "torch.Tensor",
|
||||
bias: "torch.Tensor" = None,
|
||||
group_size: int = -1,
|
||||
format: str = "NN",
|
||||
only_return_weight: bool = False,
|
||||
):
|
||||
"""
|
||||
format = TN,NN
|
||||
group_size = TN(128),NN(128, 32)
|
||||
input : bfloat16|fp16 (bs, ic)
|
||||
qweights : int32 NN: (ic, oc // 8) TN:(oc, ic // 8)
|
||||
scales : bfloat16|fp16 (ic // group_size, oc)
|
||||
zeros : int32 (ic // group_size, oc // 8)
|
||||
bias : bfloat16|fp16 (oc, )
|
||||
output : bfloat16|fp16 (bs, oc)
|
||||
"""
|
||||
|
||||
def unpack_tensor(x, pack_num=8, order_map=None):
|
||||
if order_map is None:
|
||||
order_map = [0, 1, 2, 3, 4, 5, 6, 7]
|
||||
unit = 32 // pack_num
|
||||
rows, cols = x.shape
|
||||
res = torch.zeros((rows, cols * pack_num), dtype=torch.int32, device=x.device)
|
||||
for col in range(cols):
|
||||
for k in range(pack_num):
|
||||
res[:, col * pack_num + order_map[k]] = (x[:, col] >> (unit * k)) & 0xF
|
||||
return res
|
||||
|
||||
scales = scales.t().contiguous()
|
||||
if format == "NN":
|
||||
zeros = unpack_tensor(zeros, order_map=[0, 2, 4, 6, 1, 3, 5, 7])
|
||||
zeros = zeros.t().contiguous()
|
||||
qweights = unpack_tensor(qweights, order_map=[0, 2, 4, 6, 1, 3, 5, 7])
|
||||
qweights = qweights.t().contiguous()
|
||||
else:
|
||||
zeros = unpack_tensor(zeros)
|
||||
zeros = zeros.t().contiguous()
|
||||
qweights = unpack_tensor(qweights)
|
||||
output_dim, input_dim = qweights.shape
|
||||
qweights = qweights.view(output_dim, input_dim // group_size, group_size)
|
||||
zeros = zeros.view(output_dim, input_dim // group_size, 1)
|
||||
scales = scales.view(output_dim, input_dim // group_size, 1)
|
||||
|
||||
qweights = (qweights - zeros) * scales
|
||||
qweights = qweights.view(output_dim, input_dim)
|
||||
if only_return_weight:
|
||||
return qweights
|
||||
output = torch.nn.functional.linear(inputs, qweights.to(inputs.dtype))
|
||||
return output, qweights
|
||||
|
||||
|
||||
def wui4a16_gemm(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
zeros: "torch.Tensor",
|
||||
bias: "torch.Tensor" = None,
|
||||
group_size: int = 128,
|
||||
format: str = "NN",
|
||||
):
|
||||
output_shape = inputs.shape[:-1] + (scales.shape[1],)
|
||||
|
||||
output = ops.infer.wui4a16_gemm(
|
||||
inputs, qweights, scales, zeros, bias, group_size, format
|
||||
)
|
||||
return output.view(output_shape)
|
||||
|
||||
|
||||
def wui4a16_gemv(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
zeros: "torch.Tensor",
|
||||
bias: "torch.Tensor" = None,
|
||||
group_size: int = 128,
|
||||
format: str = "NN",
|
||||
):
|
||||
output_shape = inputs.shape[:-1] + (scales.shape[1],)
|
||||
|
||||
output = ops.infer.wui4a16_gemv(
|
||||
inputs, qweights, scales, zeros, bias, group_size, format
|
||||
)
|
||||
return output.view(output_shape)
|
||||
|
||||
|
||||
def wui4a16(
|
||||
inputs: "torch.Tensor",
|
||||
qweights: "torch.Tensor",
|
||||
scales: "torch.Tensor",
|
||||
zeros: "torch.Tensor",
|
||||
bias: "torch.Tensor" = None,
|
||||
group_size: int = 128,
|
||||
format: str = "NN",
|
||||
):
|
||||
"""
|
||||
format = TN,NN
|
||||
group_size = TN(128),NN(128, 32)
|
||||
input : bfloat16|fp16 (bs, ic)
|
||||
qweights : int32 NN: (ic, oc // 8) TN:(oc, ic // 8)
|
||||
scales : bfloat16|fp16 (ic // group_size, oc)
|
||||
zeros : int32 (ic // group_size, oc // 8)
|
||||
bias : bfloat16|fp16 (oc, )
|
||||
output : bfloat16|fp16 (bs, oc)
|
||||
支持条件 : NN: oc % 8 == 0 && ic % group_size == 0 && ic % 2 == 0
|
||||
TN: oc % 2 == 0 && ic % group_size == 0
|
||||
"""
|
||||
batch = inputs.numel() // inputs.shape[-1]
|
||||
if batch <= 1:
|
||||
return wui4a16_gemv(
|
||||
inputs=inputs,
|
||||
qweights=qweights,
|
||||
scales=scales,
|
||||
zeros=zeros,
|
||||
bias=bias,
|
||||
group_size=group_size,
|
||||
format=format,
|
||||
)
|
||||
else:
|
||||
return wui4a16_gemm(
|
||||
inputs=inputs,
|
||||
qweights=qweights,
|
||||
scales=scales,
|
||||
zeros=zeros,
|
||||
bias=bias,
|
||||
group_size=group_size,
|
||||
format=format,
|
||||
)
|
||||
Reference in New Issue
Block a user