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

ixformer_sdk/inference/distributed/__init__.py

@@ -0,0 +1 @@
+from .mpi_utils import *

ixformer_sdk/inference/distributed/mpi_utils.py

@@ -0,0 +1,21 @@
+import os
+
+from mpi4py import MPI
+
+
+def get_world_size(comm=None):
+    if comm is None:
+        comm = MPI.COMM_WORLD
+
+    return comm.Get_size()
+
+
+def get_local_rank(comm=None):
+    return int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])
+
+
+def get_rank(comm=None):
+    if comm is None:
+        comm = MPI.COMM_WORLD
+
+    return comm.Get_rank()

ixformer_sdk/inference/functions/__init__.py

@@ -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 *

ixformer_sdk/inference/functions/act_and_mul.py

@@ -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

ixformer_sdk/inference/functions/act_bias_mm.py

@@ -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

ixformer_sdk/inference/functions/add.py

@@ -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

ixformer_sdk/inference/functions/bert.py

@@ -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

ixformer_sdk/inference/functions/bnb_dequant.py

@@ -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)

ixformer_sdk/inference/functions/bnb_double_quant.py

@@ -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

ixformer_sdk/inference/functions/bnb_mm_dequant.py

@@ -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

ixformer_sdk/inference/functions/bnb_qgemm.py

@@ -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)

ixformer_sdk/inference/functions/bnb_quant.py

@@ -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)

ixformer_sdk/inference/functions/bnb_rowcol_absmax.py

@@ -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)

ixformer_sdk/inference/functions/conv2d.py

@@ -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

ixformer_sdk/inference/functions/cross_entropy_loss.py

@@ -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

ixformer_sdk/inference/functions/flash_attn.py

@@ -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

ixformer_sdk/inference/functions/flash_attn_lib.py

@@ -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)

ixformer_sdk/inference/functions/fused_rope.py

@@ -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

ixformer_sdk/inference/functions/gemv.py

@@ -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

ixformer_sdk/inference/functions/groupnorm.py

@@ -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

ixformer_sdk/inference/functions/i8w8o32.py

@@ -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

ixformer_sdk/inference/functions/layernorm.py

@@ -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

ixformer_sdk/inference/functions/lightllm.py

@@ -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

ixformer_sdk/inference/functions/linalg.py

@@ -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)

ixformer_sdk/inference/functions/linear.py

@@ -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

ixformer_sdk/inference/functions/lmdeploy.py

@@ -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

ixformer_sdk/inference/functions/marlin.py

@@ -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

ixformer_sdk/inference/functions/matmul.py

@@ -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
+    )

ixformer_sdk/inference/functions/mla_fused.py

@@ -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

ixformer_sdk/inference/functions/mm.py

@@ -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

ixformer_sdk/inference/functions/overlap_comm.py

@@ -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

ixformer_sdk/inference/functions/paged_attention.py

@@ -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()

ixformer_sdk/inference/functions/quantized_linear.py

@@ -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()

ixformer_sdk/inference/functions/residual_bias.py

@@ -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

ixformer_sdk/inference/functions/rms_norm.py

@@ -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

ixformer_sdk/inference/functions/scaled_dot_product_attention.py

@@ -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)

ixformer_sdk/inference/functions/smoothquant.py

@@ -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

ixformer_sdk/inference/functions/softmax.py

@@ -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

ixformer_sdk/inference/functions/store_kv_cache.py

@@ -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,
+    )

ixformer_sdk/inference/functions/t5.py

@@ -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

ixformer_sdk/inference/functions/tgi.py

@@ -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),
+    )

ixformer_sdk/inference/functions/w8a16.py

@@ -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)

ixformer_sdk/inference/functions/w8a8.py

@@ -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,
+        )

ixformer_sdk/inference/functions/wi4a16.py

@@ -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,
+        )

ixformer_sdk/inference/functions/wui4a16.py

@@ -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,
+        )

ixformer_sdk/inference/models/clip/__init__.py

@@ -0,0 +1 @@
+from .modeling_clip import CLIPModel

ixformer_sdk/inference/models/clip/configuration_clip.py

@@ -0,0 +1,503 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" CLIP model configuration"""
+
+import copy
+import os
+from collections import OrderedDict
+from typing import TYPE_CHECKING, Any, Mapping, Optional, Union
+
+if TYPE_CHECKING:
+    from transformers.processing_utils import ProcessorMixin
+    from transformers.utils import TensorType
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.onnx import OnnxConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "openai/clip-vit-base-patch32": "https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/config.json",
+    # See all CLIP models at https://huggingface.co/models?filter=clip
+}
+
+
+class CLIPTextConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`CLIPTextModel`]. It is used to instantiate a CLIP
+    text encoder according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the text encoder of the CLIP
+    [openai/clip-vit-base-patch32](https://huggingface.co/openai/clip-vit-base-patch32) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 49408):
+            Vocabulary size of the CLIP text model. Defines the number of different tokens that can be represented by
+            the `inputs_ids` passed when calling [`CLIPModel`].
+        hidden_size (`int`, *optional*, defaults to 512):
+            Dimensionality of the encoder layers and the pooler layer.
+        intermediate_size (`int`, *optional*, defaults to 2048):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        num_hidden_layers (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 8):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        max_position_embeddings (`int`, *optional*, defaults to 77):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        hidden_act (`str` or `function`, *optional*, defaults to `"quick_gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"selu"` and `"gelu_new"` `"quick_gelu"` are supported.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-5):
+            The epsilon used by the layer normalization layers.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        initializer_factor (`float`, *optional*, defaults to 1):
+            A factor for initializing all weight matrices (should be kept to 1, used internally for initialization
+            testing).
+
+    Example:
+
+    ```python
+    >>> from transformers import CLIPTextConfig, CLIPTextModel
+
+    >>> # Initializing a CLIPTextConfig with openai/clip-vit-base-patch32 style configuration
+    >>> configuration = CLIPTextConfig()
+
+    >>> # Initializing a CLIPTextModel (with random weights) from the openai/clip-vit-base-patch32 style configuration
+    >>> model = CLIPTextModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+    model_type = "clip_text_model"
+
+    def __init__(
+        self,
+        vocab_size=49408,
+        hidden_size=512,
+        intermediate_size=2048,
+        projection_dim=512,
+        num_hidden_layers=12,
+        num_attention_heads=8,
+        max_position_embeddings=77,
+        hidden_act="quick_gelu",
+        layer_norm_eps=1e-5,
+        attention_dropout=0.0,
+        initializer_range=0.02,
+        initializer_factor=1.0,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=2,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.projection_dim = projection_dim
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.initializer_factor = initializer_factor
+        self.attention_dropout = attention_dropout
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
+    ) -> "PretrainedConfig":
+        config_dict, kwargs = cls.get_config_dict(
+            pretrained_model_name_or_path, **kwargs
+        )
+
+        # get the text config dict if we are loading from CLIPConfig
+        if config_dict.get("model_type") == "clip":
+            config_dict = config_dict["text_config"]
+
+        if (
+            "model_type" in config_dict
+            and hasattr(cls, "model_type")
+            and config_dict["model_type"] != cls.model_type
+        ):
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class CLIPVisionConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`CLIPVisionModel`]. It is used to instantiate a
+    CLIP vision encoder according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the vision encoder of the CLIP
+    [openai/clip-vit-base-patch32](https://huggingface.co/openai/clip-vit-base-patch32) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        hidden_size (`int`, *optional*, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        intermediate_size (`int`, *optional*, defaults to 3072):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        num_hidden_layers (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        image_size (`int`, *optional*, defaults to 224):
+            The size (resolution) of each image.
+        patch_size (`int`, *optional*, defaults to 32):
+            The size (resolution) of each patch.
+        hidden_act (`str` or `function`, *optional*, defaults to `"quick_gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"selu"` and `"gelu_new"` ``"quick_gelu"` are supported.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-5):
+            The epsilon used by the layer normalization layers.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        initializer_factor (`float`, *optional*, defaults to 1):
+            A factor for initializing all weight matrices (should be kept to 1, used internally for initialization
+            testing).
+
+    Example:
+
+    ```python
+    >>> from transformers import CLIPVisionConfig, CLIPVisionModel
+
+    >>> # Initializing a CLIPVisionConfig with openai/clip-vit-base-patch32 style configuration
+    >>> configuration = CLIPVisionConfig()
+
+    >>> # Initializing a CLIPVisionModel (with random weights) from the openai/clip-vit-base-patch32 style configuration
+    >>> model = CLIPVisionModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "clip_vision_model"
+
+    def __init__(
+        self,
+        hidden_size=768,
+        intermediate_size=3072,
+        projection_dim=512,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        num_channels=3,
+        image_size=224,
+        patch_size=32,
+        hidden_act="quick_gelu",
+        layer_norm_eps=1e-5,
+        attention_dropout=0.0,
+        initializer_range=0.02,
+        initializer_factor=1.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.projection_dim = projection_dim
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.initializer_range = initializer_range
+        self.initializer_factor = initializer_factor
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
+    ) -> "PretrainedConfig":
+        config_dict, kwargs = cls.get_config_dict(
+            pretrained_model_name_or_path, **kwargs
+        )
+
+        # get the vision config dict if we are loading from CLIPConfig
+        if config_dict.get("model_type") == "clip":
+            config_dict = config_dict["vision_config"]
+
+        if (
+            "model_type" in config_dict
+            and hasattr(cls, "model_type")
+            and config_dict["model_type"] != cls.model_type
+        ):
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
+            )
+
+        return cls.from_dict(config_dict, **kwargs)
+
+
+class CLIPConfig(PretrainedConfig):
+    r"""
+    [`CLIPConfig`] is the configuration class to store the configuration of a [`CLIPModel`]. It is used to instantiate
+    a CLIP model according to the specified arguments, defining the text model and vision model configs. Instantiating
+    a configuration with the defaults will yield a similar configuration to that of the CLIP
+    [openai/clip-vit-base-patch32](https://huggingface.co/openai/clip-vit-base-patch32) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        text_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize [`CLIPTextConfig`].
+        vision_config (`dict`, *optional*):
+            Dictionary of configuration options used to initialize [`CLIPVisionConfig`].
+        projection_dim (`int`, *optional*, defaults to 512):
+            Dimentionality of text and vision projection layers.
+        logit_scale_init_value (`float`, *optional*, defaults to 2.6592):
+            The inital value of the *logit_scale* paramter. Default is used as per the original CLIP implementation.
+        kwargs (*optional*):
+            Dictionary of keyword arguments.
+
+    Example:
+
+    ```python
+    >>> from transformers import CLIPConfig, CLIPModel
+
+    >>> # Initializing a CLIPConfig with openai/clip-vit-base-patch32 style configuration
+    >>> configuration = CLIPConfig()
+
+    >>> # Initializing a CLIPModel (with random weights) from the openai/clip-vit-base-patch32 style configuration
+    >>> model = CLIPModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+
+    >>> # We can also initialize a CLIPConfig from a CLIPTextConfig and a CLIPVisionConfig
+    >>> from transformers import CLIPTextConfig, CLIPVisionConfig
+
+    >>> # Initializing a CLIPText and CLIPVision configuration
+    >>> config_text = CLIPTextConfig()
+    >>> config_vision = CLIPVisionConfig()
+
+    >>> config = CLIPConfig.from_text_vision_configs(config_text, config_vision)
+    ```"""
+
+    model_type = "clip"
+    is_composition = True
+
+    def __init__(
+        self,
+        text_config=None,
+        vision_config=None,
+        projection_dim=512,
+        logit_scale_init_value=2.6592,
+        **kwargs,
+    ):
+        # If `_config_dict` exist, we use them for the backward compatibility.
+        # We pop out these 2 attributes before calling `super().__init__` to avoid them being saved (which causes a lot
+        # of confusion!).
+        text_config_dict = kwargs.pop("text_config_dict", None)
+        vision_config_dict = kwargs.pop("vision_config_dict", None)
+
+        super().__init__(**kwargs)
+
+        # Instead of simply assigning `[text|vision]_config_dict` to `[text|vision]_config`, we use the values in
+        # `[text|vision]_config_dict` to update the values in `[text|vision]_config`. The values should be same in most
+        # cases, but we don't want to break anything regarding `_config_dict` that existed before commit `8827e1b2`.
+        if text_config_dict is not None:
+            if text_config is None:
+                text_config = {}
+
+            # This is the complete result when using `text_config_dict`.
+            _text_config_dict = CLIPTextConfig(**text_config_dict).to_dict()
+
+            # Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
+            for key, value in _text_config_dict.items():
+                if (
+                    key in text_config
+                    and value != text_config[key]
+                    and key not in ["transformers_version"]
+                ):
+                    # If specified in `text_config_dict`
+                    if key in text_config_dict:
+                        message = (
+                            f"`{key}` is found in both `text_config_dict` and `text_config` but with different values. "
+                            f'The value `text_config_dict["{key}"]` will be used instead.'
+                        )
+                    # If inferred from default argument values (just to be super careful)
+                    else:
+                        message = (
+                            f"`text_config_dict` is provided which will be used to initialize `CLIPTextConfig`. The "
+                            f'value `text_config["{key}"]` will be overriden.'
+                        )
+                    logger.warning(message)
+
+            # Update all values in `text_config` with the ones in `_text_config_dict`.
+            text_config.update(_text_config_dict)
+
+        if vision_config_dict is not None:
+            if vision_config is None:
+                vision_config = {}
+
+            # This is the complete result when using `vision_config_dict`.
+            _vision_config_dict = CLIPVisionConfig(**vision_config_dict).to_dict()
+            # convert keys to string instead of integer
+            if "id2label" in _vision_config_dict:
+                _vision_config_dict["id2label"] = {
+                    str(key): value
+                    for key, value in _vision_config_dict["id2label"].items()
+                }
+
+            # Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
+            for key, value in _vision_config_dict.items():
+                if (
+                    key in vision_config
+                    and value != vision_config[key]
+                    and key not in ["transformers_version"]
+                ):
+                    # If specified in `vision_config_dict`
+                    if key in vision_config_dict:
+                        message = (
+                            f"`{key}` is found in both `vision_config_dict` and `vision_config` but with different "
+                            f'values. The value `vision_config_dict["{key}"]` will be used instead.'
+                        )
+                    # If inferred from default argument values (just to be super careful)
+                    else:
+                        message = (
+                            f"`vision_config_dict` is provided which will be used to initialize `CLIPVisionConfig`. "
+                            f'The value `vision_config["{key}"]` will be overriden.'
+                        )
+                    logger.warning(message)
+
+            # Update all values in `vision_config` with the ones in `_vision_config_dict`.
+            vision_config.update(_vision_config_dict)
+
+        if text_config is None:
+            text_config = {}
+            logger.info(
+                "`text_config` is `None`. Initializing the `CLIPTextConfig` with default values."
+            )
+
+        if vision_config is None:
+            vision_config = {}
+            logger.info(
+                "`vision_config` is `None`. initializing the `CLIPVisionConfig` with default values."
+            )
+
+        self.text_config = CLIPTextConfig(**text_config)
+        self.vision_config = CLIPVisionConfig(**vision_config)
+
+        self.projection_dim = projection_dim
+        self.logit_scale_init_value = logit_scale_init_value
+        self.initializer_factor = 1.0
+
+    @classmethod
+    def from_text_vision_configs(
+        cls, text_config: CLIPTextConfig, vision_config: CLIPVisionConfig, **kwargs
+    ):
+        r"""
+        Instantiate a [`CLIPConfig`] (or a derived class) from clip text model configuration and clip vision model
+        configuration.
+
+        Returns:
+            [`CLIPConfig`]: An instance of a configuration object
+        """
+
+        return cls(
+            text_config=text_config.to_dict(),
+            vision_config=vision_config.to_dict(),
+            **kwargs,
+        )
+
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
+
+        Returns:
+            `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+        """
+        output = copy.deepcopy(self.__dict__)
+        output["text_config"] = self.text_config.to_dict()
+        output["vision_config"] = self.vision_config.to_dict()
+        output["model_type"] = self.__class__.model_type
+        return output
+
+
+class CLIPOnnxConfig(OnnxConfig):
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        return OrderedDict(
+            [
+                ("input_ids", {0: "batch", 1: "sequence"}),
+                (
+                    "pixel_values",
+                    {0: "batch", 1: "num_channels", 2: "height", 3: "width"},
+                ),
+                ("attention_mask", {0: "batch", 1: "sequence"}),
+            ]
+        )
+
+    @property
+    def outputs(self) -> Mapping[str, Mapping[int, str]]:
+        return OrderedDict(
+            [
+                ("logits_per_image", {0: "batch"}),
+                ("logits_per_text", {0: "batch"}),
+                ("text_embeds", {0: "batch"}),
+                ("image_embeds", {0: "batch"}),
+            ]
+        )
+
+    @property
+    def atol_for_validation(self) -> float:
+        return 1e-4
+
+    def generate_dummy_inputs(
+        self,
+        processor: "ProcessorMixin",
+        batch_size: int = -1,
+        seq_length: int = -1,
+        framework: Optional["TensorType"] = None,
+    ) -> Mapping[str, Any]:
+        text_input_dict = super().generate_dummy_inputs(
+            processor.tokenizer,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            framework=framework,
+        )
+        image_input_dict = super().generate_dummy_inputs(
+            processor.feature_extractor, batch_size=batch_size, framework=framework
+        )
+        return {**text_input_dict, **image_input_dict}
+
+    @property
+    def default_onnx_opset(self) -> int:
+        return 14

ixformer_sdk/inference/models/codeshell_7b_chat/__init__.py

@@ -0,0 +1 @@
+from .modeling_codeshell import CodeShellForCausalLM

ixformer_sdk/inference/models/codeshell_7b_chat/configuration_codeshell.py

@@ -0,0 +1,153 @@
+# coding=utf-8
+# Copyright 2023 WisdomShell Inc. All Rights Reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# This code is based on Bigcode's GPTBigCode configuration. It has been modified from
+# its original forms to accommodate minor architectural differences compared to
+# GPTBigCode Configuration that trained the model.
+
+# coding=utf-8
+# Copyright 2023 The BigCode team and HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" CodeShell configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class CodeShellConfig(PretrainedConfig):
+    """
+    This is the configuration class to store the configuration of a [`CodeShellModel`]. It is used to instantiate a
+    CodeShell model according to the specified arguments, defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 50257):
+            Vocabulary size of the GPT-2 model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`CodeShellModel`].
+        n_positions (`int`, *optional*, defaults to 1024):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        n_embd (`int`, *optional*, defaults to 768):
+            Dimensionality of the embeddings and hidden states.
+        n_layer (`int`, *optional*, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        n_head (`int`, *optional*, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        n_inner (`int`, *optional*, defaults to None):
+            Dimensionality of the inner feed-forward layers. `None` will set it to 4 times n_embd
+        activation_function (`str`, *optional*, defaults to `"gelu_pytorch_tanh"`):
+            Activation function, to be selected in the list `["relu", "silu", "gelu", "tanh", "gelu_new",
+            "gelu_pytorch_tanh"]`.
+        resid_pdrop (`float`, *optional*, defaults to 0.1):
+            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+        embd_pdrop (`float`, *optional*, defaults to 0.1):
+            The dropout ratio for the embeddings.
+        attn_pdrop (`float`, *optional*, defaults to 0.1):
+            The dropout ratio for the attention.
+        layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
+            The epsilon to use in the layer normalization layers.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        scale_attn_weights (`bool`, *optional*, defaults to `True`):
+            Scale attention weights by dividing by sqrt(hidden_size)..
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models).
+        attention_softmax_in_fp32 (`bool`, *optional*, defaults to `True`):
+            Whether to call the fused softmax in float32.
+        scale_attention_softmax_in_fp32 (`bool`, *optional*, defaults to `True`):
+            Whether to scale the attention softmax in float32.
+        attention_type (`bool`, *optional*, defaults to `True`):
+            Whether to use Multi-Query Attion (`True`) or Multi-Head Attention (`False`).
+    """
+
+    model_type = "codeshell"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {
+        "hidden_size": "n_embd",
+        "max_position_embeddings": "n_positions",
+        "num_attention_heads": "n_head",
+        "num_hidden_layers": "n_layer",
+    }
+
+    def __init__(
+        self,
+        vocab_size=70144,
+        n_positions=8192,
+        n_embd=4096,
+        n_layer=42,
+        n_head=32,
+        n_inner=None,
+        activation_function="gelu_pytorch_tanh",
+        resid_pdrop=0.1,
+        embd_pdrop=0.1,
+        attn_pdrop=0.1,
+        layer_norm_epsilon=1e-5,
+        initializer_range=0.02,
+        scale_attn_weights=True,
+        use_cache=True,
+        bos_token_id=70000,
+        eos_token_id=70000,
+        attention_softmax_in_fp32=True,
+        scale_attention_softmax_in_fp32=True,
+        group_query_attention=True,
+        num_query_groups=1,
+        position_embedding_type="learned_absolute",
+        rope_scaling=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_positions = n_positions
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_inner = n_inner
+        self.activation_function = activation_function
+        self.resid_pdrop = resid_pdrop
+        self.embd_pdrop = embd_pdrop
+        self.attn_pdrop = attn_pdrop
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range
+        self.scale_attn_weights = scale_attn_weights
+        self.use_cache = use_cache
+        self.attention_softmax_in_fp32 = attention_softmax_in_fp32
+        self.scale_attention_softmax_in_fp32 = scale_attention_softmax_in_fp32
+        self.group_query_attention = group_query_attention
+        self.num_query_groups = num_query_groups
+        self.position_embedding_type = position_embedding_type
+        self.rope_scaling = rope_scaling
+        assert self.position_embedding_type in [
+            "learned_absolute",
+            "rope",
+        ], "position_embedding_type must be one of ['learned_absolute', 'rope']"
+
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+
+        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)

ixformer_sdk/inference/models/codeshell_7b_chat/modeling_codeshell_ixformer.py

@@ -0,0 +1,99 @@
+import math
+
+import ixformer.functions as ixf_F
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+def mha(query, key, value, attention_mask):
+    if attention_mask is None and query.shape[2] == key.shape[2]:
+        context_layer = ixf_F.scaled_dot_product_attention(
+            query.contiguous(), key.contiguous(), value.contiguous(), is_causal=True
+        )
+    else:
+        # if attention_mask is not None:
+        #     # attention_mask = attention_mask
+        #     attention_mask = (~attention_mask).cuda().float()*(-10000)
+        context_layer = ixf_F.scaled_dot_product_attention(
+            query.contiguous(), key.contiguous(), value.contiguous(), attention_mask
+        )
+
+    # context_layer = context_layer.transpose(1, 2).contiguous()
+    # res_shape = list(context_layer.shape)
+    # res_shape = res_shape[:2] + [-1]
+    # context_layer = context_layer.view(*res_shape)
+    return context_layer
+    # batch_size, head_num, seq_len, head_dim = query.shape
+    # src_len = query.shape[-2]
+    # tgt_len = key.shape[-2]
+
+    # if attention_mask is None and src_len == tgt_len:
+    #     attention_mask = ~torch.tril(torch.ones([src_len, tgt_len])).bool()
+    # elif attention_mask is None:
+    #     attention_mask = torch.zeros([src_len, tgt_len])
+    # attention_mask = attention_mask.cuda().int()
+
+    # attention_scores = ixf_F.act_bias_mm(
+    #     query, key, scale=1 / math.sqrt(head_dim), trans_format="TN"
+    # )
+    # # softmax
+    # # if tgt_len > 2048:
+    # #     if not (attention_mask == 0).all():
+    # #         attention_scores.masked_fill_(attention_mask.bool(), -10000.0)
+    # #     dtype = attention_scores.dtype
+    # #     attention_probs = F.softmax(attention_scores.float(), dim=-1)
+    # #     attention_probs = attention_probs.type(dtype)
+    # # else:
+    # #     raise NotImplementedError()
+    # attention_probs = ixf_F.attention_masked_softmax(
+    #     attention_scores, attention_mask.int()
+    # )
+    # # s * v
+    # # batch_size,head_num,seq_len,head_dim
+    # context_layer = ixf_F.act_bias_mm(
+    #     attention_probs, value, trans_format="NN")
+    # context_layer = context_layer.transpose(1, 2).contiguous()
+    # context_layer = context_layer.view(
+    #     batch_size, seq_len, head_num * head_dim)
+
+
+def mlp(mlp_input, ff1_weight, ff1_bias, ff2_weight):
+    input_shape = list(mlp_input.shape)
+    mlp_input = mlp_input.view(-1, input_shape[-1])
+    mlp_output = ixf_F.act_bias_mm(
+        mlp_input, ff1_weight, ff1_bias, scale=1, act_type="gelu", trans_format="TN"
+    )
+    mlp_output = ixf_F.linear(mlp_output, ff2_weight, None)
+    input_shape[-1] = -1
+    mlp_output = mlp_output.view(*input_shape)
+    return mlp_output
+
+
+def mlp_forward(self, hidden_states):
+    # [s, b, 4hp]
+    # intermediate_parallel = self.dense_h_to_4h(hidden_states)
+    # intermediate_parallel = self.activation_func(intermediate_parallel)
+    input_shape = list(hidden_states.shape)
+    hidden_states = hidden_states.view(-1, input_shape[-1])
+    mlp_output = ixf_F.act_bias_mm(
+        hidden_states,
+        self.c_fc.weight,
+        self.c_fc.bias,
+        scale=1,
+        act_type="gelu",
+        trans_format="TN",
+    )
+    if isinstance(self.c_proj, nn.Linear):
+        output = ixf_F.linear(
+            mlp_output,
+            self.c_proj.weight,
+            self.c_proj.bias,
+        )
+    else:
+        output = self.c_proj(mlp_output)
+    output = output.view(*input_shape)
+    return output

ixformer_sdk/inference/overlap/__init__.py

@@ -0,0 +1 @@
+from .llama_decoder_layer_overlap import LlamaDecoderLayerOverlapProtocol, LlamaDecoderLayerOverlapDefault, create_vllm_llama_decoder_layer

ixformer_sdk/inference/overlap/fmha_oproj_allreduce_ln_gating_overlap.py

@@ -0,0 +1,333 @@
+import dataclasses
+from typing import Optional, Tuple
+
+import ixformer.distributed as ixfd
+import ixformer.functions as F
+import torch
+import torch.distributed as dist
+from ixformer.contrib.vllm_flash_attn import flash_attn_varlen_func
+from ixformer.distributed.overlap_comm import SplitOverlapComm
+
+from ixformer.core import config as ixff_config
+
+
+@dataclasses.dataclass
+class FmhaOProjAllReduceLnGatingParams:
+    # ==============================
+    # attention
+    # ==============================
+
+    # shape: [Batch * SeqLen, NumHeads / TP, HeadDim]
+    q: torch.Tensor
+
+    # shape: [Batch * SeqLen, NumHeads / TP, HeadDim]
+    k: torch.Tensor
+
+    # shape: [Batch * SeqLen, NumHeads / TP, HeadDim]
+    v: torch.Tensor
+
+    # shape [Batch + 1], dtype torch.int32. The cumulative sequence lengths
+    # of the sequences in the batch, used to index into q.
+    cu_seqlens_q: torch.Tensor
+
+    # shape: [Batch + 1], dtype torch.int32. The cumulative sequence lengths
+    # of the sequences in the batch, used to index into kv.
+    cu_seqlens_k: torch.Tensor
+
+    # Maximum query sequence length in the batch.
+    max_seqlen_q: int
+
+    # Maximum key sequence length in the batch.
+    max_seqlen_k: int
+
+    # ==============================
+    # o_proj
+    # ==============================
+
+    # shape: [HiddenSize, NumHeads * HeadDim / TP], dtype: int8
+    o_proj_weight: torch.Tensor
+
+    # shape: [HiddenSize], dtype: float32
+    o_proj_weight_scale: torch.Tensor
+
+    # shape: [HiddenSize]
+    o_proj_bias: torch.Tensor
+
+    # shape: [NumHeads * HeadDim / TP], dtype: float16 or bfloat16
+    o_proj_smooth_scale: torch.Tensor
+
+    # ==============================
+    # ln
+    # ==============================
+
+    # shape: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16
+    residual: torch.Tensor
+
+    # shape: [HiddenSize], dtype: float16 or bfloat16
+    ln_weight: torch.Tensor
+
+    # shape: [HiddenSize], dtype: float16 or bfloat16
+    ln_bias: torch.Tensor
+
+    # ==============================
+    # gating linear
+    # ==============================
+
+    # shape: [TopK, HiddenSize], dtype: float16 or bfloat16
+    gating_weight: torch.Tensor
+
+    # shape: [SeqLen, TopK], dtype: float16 or bfloat16
+    out: Optional[torch.Tensor] = None
+
+    # ==============================
+    # default parameters
+    # ==============================
+
+    # the seqlens of q for per chunk when using overlap,
+    # the parameter can be initiated by params.prepare_overlap_params(),
+    # and only need to initialize once during the model's forward.
+    cu_seqlens_q_chunks = None
+    cu_seqlens_k_chunks = None
+
+    softmax_scale: Optional[float] = None
+    ln_eps: float = 1e-5
+
+    @property
+    def batch(self):
+        return len(self.cu_seqlens_q) - 1
+
+    @property
+    def seqlen(self):
+        return self.q.shape[0]
+
+    @property
+    def topk(self):
+        return self.gating_weight.shape[0]
+
+    def prepare_overlap_params(self):
+        """compute the cu_seqlens qk of chunk when using overlap"""
+        first_chunk_size = int(self.q.shape[0] // 2)
+
+        if not hasattr(self.cu_seqlens_q, "q_chunks"):
+            first_q_chunks_cu_seqlens = self.cu_seqlens_q.clone()
+            first_q_chunks_cu_seqlens[-1] = first_chunk_size
+
+            self.cu_seqlens_q.q_chunks = [
+                first_q_chunks_cu_seqlens,
+                first_q_chunks_cu_seqlens,
+            ]
+
+        self.cu_seqlens_q_chunks = self.cu_seqlens_q.q_chunks
+
+        if not hasattr(self.cu_seqlens_k, "k_chunks"):
+            first_chunk = self.cu_seqlens_k.clone()
+            last_chunk = self.cu_seqlens_k
+            first_chunk[-1] = first_chunk_size
+            self.cu_seqlens_k.k_chunks = [first_chunk, last_chunk]
+
+        self.cu_seqlens_k_chunks = self.cu_seqlens_k.k_chunks
+
+        return self
+
+
+class FmhaOProjAllreduceLnGatingOverlap(SplitOverlapComm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if self.num_chunks != 2:
+            raise RuntimeError(
+                f"Overlap only support num_chunks == 2, but got {self.num_chunks}."
+            )
+
+        self.allreduce_end_events = [torch.cuda.Event() for _ in range(self.num_chunks)]
+
+    def start_ln_gating(self, chunk_idx):
+        compute_stream = self._compute_streams[chunk_idx % self.num_compute_streams]
+        compute_stream.wait_event(self.allreduce_end_events[chunk_idx])
+
+    def compute(self, params: FmhaOProjAllReduceLnGatingParams):
+        if params.out is None:
+            params.out = torch.empty(
+                [params.seqlen, params.topk], device="cuda", dtype=torch.float
+            )
+
+        seqlen_chunks = [int(params.q.shape[0] // 2)]
+        seqlen_chunks.append(params.q.shape[0] - seqlen_chunks[0])
+
+        q_chunks = torch.split_with_sizes(params.q, seqlen_chunks, dim=0)
+
+        ar_out_chunks = []
+        for chunk_idx in range(len(seqlen_chunks)):
+            with self.compute_stream_context(chunk_idx):
+                hidden_states = flash_attn_varlen_func(
+                    q=q_chunks[chunk_idx],
+                    k=params.k,
+                    v=params.v,
+                    cu_seqlens_q=params.cu_seqlens_q_chunks[chunk_idx],
+                    cu_seqlens_k=params.cu_seqlens_k_chunks[chunk_idx],
+                    max_seqlen_q=seqlen_chunks[chunk_idx],
+                    max_seqlen_k=seqlen_chunks[0]
+                    if chunk_idx == 0
+                    else params.max_seqlen_k,
+                    softmax_scale=params.softmax_scale,
+                    causal=True,
+                    window_size=(-1, -1),
+                    alibi_slopes=None,
+                    softcap=0,
+                )
+
+                hidden_states = hidden_states.view(hidden_states.shape[0], -1)
+                hidden_states, i_scales = F.dynamic_scaled_quant_dynamic_int8(
+                    hidden_states, params.o_proj_smooth_scale
+                )
+
+                out_chunk = F.w8a8(
+                    hidden_states,
+                    params.o_proj_weight,
+                    i_scales,
+                    params.o_proj_weight_scale,
+                    bias=params.o_proj_bias,
+                    out_dtype=params.residual.dtype,
+                    output=None,
+                    persistent=True,
+                )
+
+            self.start_comm(chunk_idx)
+
+            ixfd.all_reduce(
+                out_chunk, async_op=True, group=self.comm_group, use_comm_stream=True
+            )
+            ar_out_chunks.append(out_chunk)
+
+            self.allreduce_end_events[chunk_idx].record(self._comm_stream)
+
+        ln_out = torch.empty_like(params.residual)
+        ln_out_chunks = ln_out.chunk(2, dim=0)
+        residual_chunks = torch.split_with_sizes(params.residual, seqlen_chunks, dim=0)
+
+        if params.out is None:
+            params.out = torch.empty(
+                [params.seqlen, params.topk], dtype=params.residual.dtype, device="cuda"
+            )
+
+        out_chunks = list(torch.split_with_sizes(params.out, seqlen_chunks, dim=0))
+
+        for chunk_idx in range(len(seqlen_chunks)):
+            self.start_ln_gating(chunk_idx)
+            with self.compute_stream_context(chunk_idx):
+                ln_out_chunk, residual_chunk = F.residual_layer_norm(
+                    input=ar_out_chunks[chunk_idx],
+                    weight=params.ln_weight,
+                    bias=params.ln_bias,
+                    residual=residual_chunks[chunk_idx].reshape(
+                        ar_out_chunks[chunk_idx].shape
+                    ),
+                    eps=params.ln_eps,
+                    output=ln_out_chunks[chunk_idx],
+                )
+                if ln_out_chunk.dtype == params.gating_weight.dtype:
+                    F.linear(
+                        ln_out_chunk, params.gating_weight, output=out_chunks[chunk_idx]
+                    )
+                else:
+                    F.mixed_type_linear(
+                        ln_out_chunk, params.gating_weight, output=out_chunks[chunk_idx]
+                    )
+
+        return (
+            params.residual.reshape(params.batch, params.seqlen, -1),
+            ln_out.reshape(params.batch, params.seqlen, -1),
+            params.out,
+        )
+
+
+_fa_o_proj_allreduce_ln_gating_overlap = None
+
+
+def fmha_oproj_allreduce_ln_gating(
+    params: FmhaOProjAllReduceLnGatingParams,
+    enable_overlap: bool = False,
+    comm_group=None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    FMHA + OProjLinear + AllReduce + LayerNorm + GatingLinear
+
+    Args:
+        params: fused operator params
+        enable_overlap: whether enable overlap
+        comm_group: communication group
+    Returns:
+        Residual: shape: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16
+        HiddenStates: shape: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16
+        GatingLinearOutput: shape: [SeqLen, TopK], dtype: float16 or bfloat16
+    """
+
+    global _fa_o_proj_allreduce_ln_gating_overlap
+    if _fa_o_proj_allreduce_ln_gating_overlap is None:
+        _fa_o_proj_allreduce_ln_gating_overlap = (
+            FmhaOProjAllreduceLnGatingOverlap.dispatcher(
+                num_chunks=2, comm_group=comm_group
+            ).forward
+        )
+
+    if (
+        enable_overlap
+        and ixff_config.IXFORMER_ENABLE_OVERLAP_COMM
+        and dist.is_initialized()
+        and dist.get_world_size(comm_group) > 1
+        and params.batch == 1
+        and params.seqlen > 1
+    ):
+        if params.cu_seqlens_q_chunks is None:
+            params = params.prepare_overlap_params()
+        return _fa_o_proj_allreduce_ln_gating_overlap(params)
+
+    hidden_states = flash_attn_varlen_func(
+        q=params.q,
+        k=params.k,
+        v=params.v,
+        cu_seqlens_q=params.cu_seqlens_q,
+        cu_seqlens_k=params.cu_seqlens_k,
+        max_seqlen_q=params.max_seqlen_q,
+        max_seqlen_k=params.max_seqlen_k,
+        softmax_scale=params.softmax_scale,
+        causal=True,
+        window_size=(-1, -1),
+        alibi_slopes=None,
+        softcap=0,
+    )
+
+    input = hidden_states.view(hidden_states.shape[0], -1)
+    input, i_scales = F.dynamic_scaled_quant_smoothquant(
+        input, params.o_proj_smooth_scale
+    )
+
+    hidden_states = F.w8a8(
+        input,
+        params.o_proj_weight,
+        i_scales,
+        params.o_proj_weight_scale,
+        bias=params.o_proj_bias,
+        out_dtype=params.residual.dtype,
+        output=None,
+    )
+    ixfd.all_reduce(hidden_states, async_op=True, group=comm_group)
+
+    hidden_states, residual = F.residual_layer_norm(
+        input=hidden_states,
+        weight=params.ln_weight,
+        bias=params.ln_bias,
+        residual=params.residual.reshape(hidden_states.shape),
+        eps=params.ln_eps,
+    )
+    if hidden_states.dtype == params.gating_weight.dtype:
+        out = F.linear(hidden_states, params.gating_weight, output=params.out)
+    else:
+        out = F.mixed_type_linear(
+            hidden_states, params.gating_weight, output=params.out
+        )
+    return (
+        residual.reshape(params.batch, params.seqlen, -1),
+        hidden_states.reshape(params.batch, params.seqlen, -1),
+        out,
+    )

ixformer_sdk/inference/overlap/group_gemm_moe_reduce_sum_allreduce_overlap.py

@@ -0,0 +1,219 @@
+import dataclasses
+import math
+from contextlib import contextmanager
+from typing import List, Optional
+
+import ixformer.distributed as ixfd
+import ixformer.functions as F
+import torch
+import torch.distributed as dist
+from ixformer.distributed.overlap_comm import SplitOverlapComm
+
+from ixformer.core import config as ixff_config
+
+
+@dataclasses.dataclass
+class GroupGemmMoeReduceSumAllReduceParams:
+    # M: NumTokens * TopK
+    # K: InnerSize // TP
+    # N: HiddenSize
+    # NumTokens: M // TopK
+
+    # =================================
+    # group gemm
+    # =================================
+
+    # the top k of experts
+    topk: int
+
+    # shape: [M, K] if format[1]=="N" else [K, M], dtype: int8
+    input: torch.Tensor
+
+    # shape: [NumExperts, N, K] if format[0]=="T" else [NumExperts, K, N], dtype: int8
+    weight: torch.Tensor
+
+    # shape: [M], dtype: float32
+    i_scales: torch.Tensor
+
+    # shape: [NumExperts, N], dtype: float32
+    w_scales: torch.Tensor
+
+    # shape: [NumExperts], dtype: int32
+    tokens_per_experts: torch.Tensor
+
+    # the dtype of output, support float16 and bfloat16
+    out_dtype: torch.dtype = None
+
+    # index of dst to src, shape: [M], dtype: int32
+    dst_to_src: torch.Tensor = None
+
+    # only support TN now
+    format: str = "TN"
+
+    # =================================
+    # moe reduce sum
+    # =================================
+
+    # shape: [M // TopK, TopK], dtype: torch.float16 or torch.bfloat16
+    topk_weight: torch.Tensor = None
+
+    # shape: [M // TopK, N], dtype: torch.float16 or torch.bfloat16
+    output: torch.Tensor = None
+
+    # overlap
+    output_chunks: Optional[List[torch.Tensor]] = None
+
+    @property
+    def M(self):
+        return self.input.shape[0]
+
+    @property
+    def N(self):
+        if torch.is_tensor(self.weight):
+            return self.weight.shape[1]
+        return sum(t.shape[1] for t in self.weight)
+
+    @property
+    def K(self):
+        return self.input.shape[-1]
+
+    def prepare_overlap_params(
+        self, num_chunks: int, split_ratio: Optional[float] = None
+    ):
+        if num_chunks == 2 and split_ratio not in [0, None]:
+            return self.prepare_overla_params_with_ratio(split_ratio)
+        return self.prepare_overlap_params_with_chunks(num_chunks)
+
+    def prepare_overlap_params_with_chunks(self, num_chunks: int):
+        if torch.is_tensor(self.weight):
+            weight_chunks = torch.chunk(self.weight, num_chunks, dim=1)
+            self.weight = list(weight_chunks)
+
+        if torch.is_tensor(self.w_scales):
+            weight_scale_chunks = torch.chunk(self.w_scales, num_chunks, dim=1)
+            self.w_scales = list(weight_scale_chunks)
+
+        if self.output is None:
+            self.output = torch.empty(
+                self.M // self.topk, self.N, dtype=self.out_dtype, device="cuda"
+            )
+
+        if torch.is_tensor(self.output):
+            output_chunks = torch.chunk(self.output, num_chunks, dim=1)
+            self.output_chunks = list(output_chunks)
+
+    def prepare_overla_params_with_ratio(self, split_ratio: float):
+        N = self.N
+        n_chunks = [int(math.ceil(N * split_ratio))]
+        n_chunks.append(N - n_chunks[0])
+
+        if torch.is_tensor(self.weight):
+            weight_chunks = torch.split(self.weight, n_chunks, dim=1)
+            self.weight = list(weight_chunks)
+
+        if torch.is_tensor(self.w_scales):
+            weight_scale_chunks = torch.split(self.w_scales, n_chunks, dim=1)
+            self.w_scales = list(weight_scale_chunks)
+
+        if self.output is None:
+            self.output = torch.empty(
+                self.M // self.topk, self.N, dtype=self.out_dtype, device="cuda"
+            )
+
+        if torch.is_tensor(self.output):
+            output_chunks = torch.split(self.output, n_chunks, dim=1)
+            self.output_chunks = list(output_chunks)
+
+
+class GroupGemmMoeReduceSumAllReduceSplitNOverlap(SplitOverlapComm):
+    def compute(self, params: GroupGemmMoeReduceSumAllReduceParams):
+        for chunk_idx, (weight, weight_scale) in enumerate(
+            zip(params.weight, params.w_scales)
+        ):
+            with self.compute_stream_context(chunk_idx):
+                out = F.moe_w8a8_group_gemm(
+                    input=params.input,
+                    weight=weight,
+                    i_scales=params.i_scales,
+                    w_scales=weight_scale,
+                    output_dtype=params.out_dtype,
+                    tokens_per_experts=params.tokens_per_experts,
+                    dst_to_src=params.dst_to_src,
+                    format=params.format,
+                )
+                out = out.reshape(-1, params.topk, out.shape[-1])
+                out = F.moe_output_reduce_sum(
+                    input=out,
+                    topk_weight=params.topk_weight,
+                    output=params.output_chunks[chunk_idx],
+                )
+
+            self.start_comm(chunk_idx)
+            ixfd.all_reduce(
+                out,
+                async_op=True,
+                group=self.comm_group,
+                use_comm_stream=True,
+                algo=ixfd.AllReduceAlgo.Stride,
+            )
+
+            # if chunk_idx == 0: torch.cuda.synchronize()
+
+        return params.output
+
+
+_group_gemm_moe_reduce_sum_all_reduce_overlap = None
+
+
+def group_gemm_moe_reduce_sum_allreduce(
+    params: GroupGemmMoeReduceSumAllReduceParams,
+    enable_overlap: bool = False,
+    comm_group=None,
+    num_chunks=2,
+    split_ratio: Optional[float] = None,
+):
+    if params.output is None and params.out_dtype is None:
+        raise RuntimeError(
+            "group_gemm_moe_reduce_sum_all_reduce need out_dtype argument when output is none."
+        )
+
+    if params.out_dtype is None:
+        params.out_dtype = params.output.dtype
+
+    if (
+        enable_overlap
+        and ixff_config.IXFORMER_ENABLE_OVERLAP_COMM
+        and dist.is_initialized()
+        and dist.get_world_size(comm_group) > 1
+    ):
+        global _group_gemm_moe_reduce_sum_all_reduce_overlap
+        if _group_gemm_moe_reduce_sum_all_reduce_overlap is None:
+            _group_gemm_moe_reduce_sum_all_reduce_overlap = (
+                GroupGemmMoeReduceSumAllReduceSplitNOverlap.dispatcher(
+                    num_chunks=num_chunks, comm_group=comm_group
+                ).forward
+            )
+
+        params.prepare_overlap_params(num_chunks=num_chunks, split_ratio=split_ratio)
+        return _group_gemm_moe_reduce_sum_all_reduce_overlap(params)
+
+    out = F.moe_w8a8_group_gemm(
+        input=params.input,
+        weight=params.weight,
+        i_scales=params.i_scales,
+        w_scales=params.w_scales,
+        output_dtype=params.out_dtype,
+        tokens_per_experts=params.tokens_per_experts,
+        dst_to_src=params.dst_to_src,
+        format=params.format,
+    )
+
+    out = out.reshape(-1, params.topk, out.shape[-1])
+    out = F.moe_output_reduce_sum(
+        input=out, topk_weight=params.topk_weight, output=params.output
+    )
+
+    if dist.is_initialized() and dist.get_world_size(comm_group) > 1:
+        ixfd.all_reduce(out, group=comm_group, async_op=True)
+
+    return out

ixformer_sdk/inference/overlap/linear_mlp_overlap_comm.py

@@ -0,0 +1,305 @@
+from contextlib import nullcontext
+from typing import List
+
+import torch.cuda
+
+from ...distributed import _distributed as ixfd
+from ...distributed import overlap_comm as base_overlap_comm
+from ...distributed.overlap_comm import GemmAllReduceSplitOverlapComm
+from .. import overlap as overlap_base
+
+
+class LinearMLPOverlapCommHook:
+    def on_mlp_linear2_finished(
+        self,
+        overlap_comm: "LinearMLPOverlapComm",
+        num_chunks,
+        chunk_idx,
+        hidden_states_chunk,
+        residual_chunk,
+    ):
+        pass
+
+    def on_mlp_finished(
+        self,
+        overlap_comm: "LinearMLPOverlapComm",
+        hidden_states_chunks,
+        residual_chunks,
+    ):
+        pass
+
+
+class LinearMLPOverlapComm(GemmAllReduceSplitOverlapComm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(num_compute_streams=1, *args, **kwargs)
+
+        self._mlp_linear1_start_events: List[torch.cuda.Event] = [
+            torch.cuda.Event() for _ in range(self.num_chunks)
+        ]
+        self._mlp_linear1_end_events: List[torch.cuda.Event] = [
+            torch.cuda.Event() for _ in range(self.num_chunks)
+        ]
+
+        self._mlp_linear1_stream: torch.cuda.Stream = torch.cuda.Stream()
+
+    def stop_linear_comm(self, chunk_idx):
+        event = self._mlp_linear1_start_events[chunk_idx]
+        event.record(self._comm_stream)
+
+    def start_mlp_linear1(self, chunk_idx):
+        self._mlp_linear1_stream.wait_event(self._mlp_linear1_start_events[chunk_idx])
+
+    def stop_mlp_linear1(self, chunk_idx):
+        event = self._mlp_linear1_end_events[chunk_idx]
+        event.record(self._mlp_linear1_stream)
+
+    def start_mlp_linear2(self, chunk_idx):
+        compute_stream = self._compute_streams[chunk_idx % self.num_compute_streams]
+        compute_stream.wait_event(self._mlp_linear1_end_events[chunk_idx])
+
+    def compute(
+        self,
+        protocol: "overlap_base.LlamaDecoderLayerOverlapDefault",
+        attn_output,
+        residual,
+        *,
+        mlp_linear2_finished_callback=None,
+        mlp_finished_callback=None,
+    ):
+        """ """
+        attn_output_shape = attn_output.shape
+        residual_shape = None if residual is None else residual.shape
+
+        is_update_shape = attn_output.ndim > 2
+        batch = 1
+        if attn_output.ndim == 2:
+            seqlen = attn_output_shape[0]
+        else:
+            batch = attn_output_shape[0]
+            seqlen = attn_output_shape[1]
+
+        parallel_dims = batch * seqlen
+
+        if is_update_shape:
+            attn_output = attn_output.reshape(parallel_dims, -1)
+            if residual is not None:
+                residual = residual.reshape(-1, residual_shape[-1])
+
+        attn_output_chunks, residual_chunks = protocol.split_mlp_inputs(
+            attn_output, residual, self.num_chunks
+        )
+
+        out = protocol.create_mlp_output()
+        out_chunks = protocol.split_mlp_output(out, self.num_chunks)
+
+        res_chunks = []
+
+        # 1. output project linear
+        for chunk_idx, (attn_output_chunk, residual_chunk) in enumerate(
+            zip(attn_output_chunks, residual_chunks)
+        ):
+            with self.compute_stream_context(chunk_idx):
+                hidden_states = protocol.attn_output_proj_linear(
+                    self.num_chunks,
+                    chunk_idx,
+                    attn_output_chunk,
+                    use_limited_gemm=chunk_idx != 0,
+                )
+
+            self.start_comm(chunk_idx)
+            ixfd.all_reduce(
+                hidden_states,
+                async_op=True,
+                group=self.comm_group,
+                use_comm_stream=True,
+            )
+            self.stop_linear_comm(chunk_idx)
+            attn_output_chunks[chunk_idx] = hidden_states
+
+        # 2. ln, mlp_linear1 and act
+        for chunk_idx, (hidden_states, residual_chunk) in enumerate(
+            zip(attn_output_chunks, residual_chunks)
+        ):
+            self.start_mlp_linear1(chunk_idx)
+            with self.stream_context(self._mlp_linear1_stream):
+                (
+                    hidden_states,
+                    residual_chunk,
+                ) = protocol.attn_output_proj_linear_layer_norm(
+                    self.num_chunks, chunk_idx, hidden_states, residual_chunk
+                )
+
+                hidden_states = protocol.mlp_linear1(
+                    self.num_chunks, chunk_idx, hidden_states, use_limited_gemm=True
+                )
+                hidden_states = protocol.mlp_activation(hidden_states)
+
+                attn_output_chunks[chunk_idx] = hidden_states
+                res_chunks.append(residual_chunk)
+
+            self.stop_mlp_linear1(chunk_idx)
+
+        # 3. mlp_linear2
+        for chunk_idx, hidden_states in enumerate(attn_output_chunks):
+            self.start_mlp_linear2(chunk_idx)
+            with self.compute_stream_context(chunk_idx):
+                hidden_states = protocol.mlp_linear2(
+                    self.num_chunks,
+                    chunk_idx,
+                    hidden_states,
+                    out=out_chunks[chunk_idx],
+                    use_limited_gemm=True,
+                )
+
+            self.start_comm(chunk_idx)
+            ixfd.all_reduce(
+                hidden_states,
+                async_op=True,
+                group=self.comm_group,
+                use_comm_stream=True,
+            )
+
+            if mlp_linear2_finished_callback is not None:
+                mlp_linear2_finished_callback(
+                    self,
+                    self.num_chunks,
+                    chunk_idx,
+                    hidden_states,
+                    res_chunks[chunk_idx],
+                )
+
+        if mlp_finished_callback is not None:
+            mlp_finished_callback(self, out_chunks, res_chunks)
+
+        if is_update_shape:
+            out = out.reshape(attn_output_shape)
+            if residual is not None:
+                residual = residual.reshape(residual_shape)
+
+        return out, residual
+
+    def gemm_dispatcher(
+        self,
+        chunk_idx,
+        chunk_input,
+        weight,
+        chunk_out,
+        use_limited_gemm=False,
+        user_gemm_method=None,
+        *args,
+        **kwargs,
+    ):
+        if user_gemm_method is not None and callable(user_gemm_method):
+            ctx = self.ixf_limited_gemm_ctx if use_limited_gemm else nullcontext()
+            with ctx:
+                return user_gemm_method(
+                    chunk_input, weight, out=chunk_out, *args, **kwargs
+                )
+
+        ctx = self.limited_gemm_ctx if use_limited_gemm else nullcontext()
+        with ctx:
+            return torch.matmul(chunk_input, weight.T, out=chunk_out)
+
+    @classmethod
+    def is_supported(cls, input, num_chunks, comm_group):
+        if not cls.enable():
+            return False
+
+        ndim = input.ndim
+        shape = input.shape
+
+        if ndim == 1:
+            m, k = 1, shape[0]
+        elif ndim == 2:
+            m, k = shape
+        else:
+            m, k = sum(shape[:-1]), shape[-1]
+
+        return m >= 512
+
+    @classmethod
+    def native_forward(
+        cls,
+        attn_output,
+        residual,
+        linear_weight,
+        ln_layer,
+        mlp_weight1,
+        mlp_weight2,
+        mlp_activation,
+        linear_method=None,
+        mlp_linear1_method=None,
+        mlp_linear2_method=None,
+        group=None,
+        *args,
+        **kwargs,
+    ):
+        import ixformer.functions as ixff
+
+        linear_method = linear_method or ixff.linear
+        mlp_linear1_method = mlp_linear1_method or ixff.linear
+        mlp_linear2_method = mlp_linear2_method or ixff.linear
+
+        hidden_states = linear_method(attn_output, linear_weight)
+        ixfd.all_reduce(hidden_states, async_op=True, group=group)
+
+        if ln_layer is not None:
+            hidden_states, residual = ln_layer(hidden_states, residual)
+
+        hidden_states = mlp_linear1_method(hidden_states, mlp_weight1)
+        hidden_states = mlp_activation(hidden_states)
+        hidden_states = mlp_linear2_method(hidden_states, mlp_weight2)
+        ixfd.all_reduce(hidden_states, async_op=True, group=group)
+
+        return hidden_states, residual
+
+
+_DEFAULT_OVERLAP_GROUP = None
+_DEFAULT_OVERLAP_COMM_N2 = None
+_DEFAULT_OVERLAP_COMM_N4 = None
+_DEFAULT_OVERLAP_CHUNKS = base_overlap_comm._DEFAULT_OVERLAP_CHUNKS
+
+
+def linear_mlp_overlap(
+    protocol: "overlap_base.LlamaDecoderLayerOverlapProtocol",
+    attn_output,
+    residual,
+    num_chunks=None,
+    group=None,
+    *,
+    mlp_linear2_finished_callback=None,
+    mlp_finished_callback=None,
+):
+    num_chunks = num_chunks or _DEFAULT_OVERLAP_CHUNKS
+
+    global _DEFAULT_OVERLAP_GROUP
+    global _DEFAULT_OVERLAP_COMM_N2
+    global _DEFAULT_OVERLAP_COMM_N4
+
+    if _DEFAULT_OVERLAP_GROUP is None:
+        _DEFAULT_OVERLAP_GROUP = group
+
+    if num_chunks == 2 and group == _DEFAULT_OVERLAP_GROUP:
+        if _DEFAULT_OVERLAP_COMM_N2 is None:
+            _DEFAULT_OVERLAP_COMM_N2 = LinearMLPOverlapComm.dispatcher(
+                num_chunks=num_chunks, comm_group=group
+            )
+        overlap_comm = _DEFAULT_OVERLAP_COMM_N2
+    elif num_chunks == 4 and group == _DEFAULT_OVERLAP_GROUP:
+        if _DEFAULT_OVERLAP_COMM_N4 is None:
+            _DEFAULT_OVERLAP_COMM_N4 = LinearMLPOverlapComm.dispatcher(
+                num_chunks=num_chunks, comm_group=group
+            )
+        overlap_comm = _DEFAULT_OVERLAP_COMM_N4
+    else:
+        overlap_comm = LinearMLPOverlapComm.dispatcher(
+            num_chunks=num_chunks, comm_group=group
+        )
+
+    return overlap_comm.forward(
+        protocol,
+        attn_output,
+        residual,
+        mlp_linear2_finished_callback=mlp_linear2_finished_callback,
+        mlp_finished_callback=mlp_finished_callback,
+    )

ixformer_sdk/inference/overlap/moe_reduce_allreduce_ln_linear_overlap.py

@@ -0,0 +1,190 @@
+import dataclasses
+import math
+from typing import Optional, Tuple
+
+import ixformer.distributed as ixfd
+import ixformer.functions as F
+import torch
+import torch.distributed as dist
+from ixformer.distributed.overlap_comm import SplitOverlapComm
+
+from ixformer.core import config as ixff_config
+
+
+@dataclasses.dataclass
+class MoeReduceAllReduceLnQkvLinearParams:
+    # ==============================
+    # MOE Reduce Sum
+    # ==============================
+
+    # shape: [Batch * SeqLen, TopK, HiddenSize], dtype: float16 or bfloat16
+    input: torch.Tensor
+
+    # shape: [Batch * SeqLen, TopK], dtype: float32
+    topk_weight: Optional[torch.Tensor]
+
+    # ==============================
+    # Ln
+    # ==============================
+
+    # shape: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16
+    residual: torch.Tensor
+
+    # shape: [HiddenSize], dtype: float16 or bfloat16
+    ln_weight: torch.Tensor
+
+    # shape: [HiddenSize], dtype: float16 or bfloat16
+    ln_bias: torch.Tensor
+    ln_eps: float
+
+    # ==============================
+    # QkvLinear
+    # ==============================
+
+    # shape: [(NumHeads + 2 * NumKvHeads) * HeadDim / TP, HiddenSize], dtype: float16 or bfloat16
+    qkv_weight: torch.Tensor
+
+    # shape: [(NumHeads + 2 * NumKvHeads) * HeadDim / TP]
+    qkv_weight_scale: torch.Tensor
+
+    # shape: [Batch * SeqLen, (NumHeads + 2 * NumKvHeads) * HeadDim / TP], dtype: float16 or bfloat16
+    qkv_out: torch.Tensor
+
+
+class MoeReduceSumAllReduceLnQkvLinearOverlap(SplitOverlapComm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.allreduce_end_events = [torch.cuda.Event() for _ in range(self.num_chunks)]
+
+    def start_qkv_linear(self, chunk_idx):
+        compute_stream = self._compute_streams[chunk_idx % self.num_compute_streams]
+        compute_stream.wait_event(self.allreduce_end_events[chunk_idx])
+
+    def compute(self, params: MoeReduceAllReduceLnQkvLinearParams, split_ratio=0.5):
+        input_chunk_sizes = [int(math.ceil(params.input.shape[0] * split_ratio))]
+        input_chunk_sizes.append(params.input.shape[0] - input_chunk_sizes[0])
+
+        input_chunks = list(
+            torch.split_with_sizes(params.input, input_chunk_sizes, dim=0)
+        )
+        topk_weight_chunks = list(
+            torch.split_with_sizes(params.topk_weight, input_chunk_sizes, dim=0)
+        )
+        out_chunks = []
+
+        for chunk_idx in range(len(input_chunks)):
+            with self.compute_stream_context(chunk_idx):
+                out_chunks.append(
+                    F.moe_output_reduce_sum(
+                        input_chunks[chunk_idx],
+                        topk_weight=topk_weight_chunks[chunk_idx],
+                    )
+                )
+
+            self.start_comm(chunk_idx)
+
+            ixfd.all_reduce(
+                out_chunks[chunk_idx],
+                async_op=True,
+                group=self.comm_group,
+                use_comm_stream=True,
+            )
+
+            self.allreduce_end_events[chunk_idx].record(self._comm_stream)
+
+        residual_chunk_sizes = [int(params.residual.shape[0] * split_ratio)]
+        residual_chunk_sizes.append(params.residual.shape[0] - residual_chunk_sizes[0])
+        residual_chunks = torch.split_with_sizes(
+            params.residual, residual_chunk_sizes, dim=0
+        )
+
+        qkv_out_chunk_sizes = [int(params.qkv_out.shape[0] * split_ratio)]
+        qkv_out_chunk_sizes.append(params.qkv_out.shape[0] - qkv_out_chunk_sizes[0])
+        qkv_out_chunks = torch.split_with_sizes(
+            params.qkv_out, qkv_out_chunk_sizes, dim=0
+        )
+
+        for chunk_idx in range(len(input_chunks)):
+            self.start_qkv_linear(chunk_idx)
+            with self.compute_stream_context(chunk_idx):
+                (
+                    i8_hidden_states,
+                    residual,
+                    i_scales,
+                ) = F.residual_layer_norm_dynamic_int8(
+                    input=out_chunks[chunk_idx],
+                    residual=residual_chunks[chunk_idx],
+                    weight=params.ln_weight,
+                    bias=params.ln_bias,
+                    eps=params.ln_eps,
+                )
+
+                F.w8a8(
+                    i8_hidden_states,
+                    params.qkv_weight,
+                    i_scales,
+                    params.qkv_weight_scale,
+                    output=qkv_out_chunks[chunk_idx],
+                )
+
+        return params.qkv_out, params.residual
+
+
+_moe_reduce_with_allreduce_overlap = None
+
+
+def moe_reduce_sum_allreduce_ln_qkv_linear(
+    params: MoeReduceAllReduceLnQkvLinearParams,
+    enable_overlap=False,
+    comm_group=None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    MOE Reduce Sum + AllReduce + LayerNorm + QkvLinear
+
+    Args:
+        params: fused operator params
+        enable_overlap: whether enable overlap
+        comm_group: communication group
+    Returns:
+        QkvLinearOutput: [Batch * SeqLen, (NumHeads + 2 * NumKvHeads) * HeadDim / TP], dtype: float16 or bfloat16
+        Residual: [Batch * SeqLen, HiddenSize], dtype: float16 or bfloat16
+    """
+
+    global _moe_reduce_with_allreduce_overlap
+    if _moe_reduce_with_allreduce_overlap is None:
+        _moe_reduce_with_allreduce_overlap = (
+            MoeReduceSumAllReduceLnQkvLinearOverlap.dispatcher(
+                num_chunks=2, comm_group=comm_group
+            ).forward
+        )
+
+    if (
+        enable_overlap
+        and ixff_config.IXFORMER_ENABLE_OVERLAP_COMM
+        and dist.is_initialized()
+        and dist.get_world_size(comm_group) > 1
+        and params.input.shape[0] > 1
+    ):
+        return _moe_reduce_with_allreduce_overlap(params, split_ratio=0.5)
+
+    out = F.moe_output_reduce_sum(params.input, topk_weight=params.topk_weight)
+    ixfd.all_reduce(out, async_op=True, group=comm_group)
+
+    i8_hidden_states, residual, i_scales = F.residual_layer_norm_dynamic_int8(
+        input=out,
+        residual=params.residual,
+        weight=params.ln_weight,
+        bias=params.ln_bias,
+        eps=params.ln_eps,
+    )
+
+    out = F.w8a8(
+        i8_hidden_states,
+        params.qkv_weight,
+        i_scales,
+        params.qkv_weight_scale,
+        output=params.qkv_out,
+    )
+
+    return out, residual

ixformer_sdk/inference/overlap/overlap_comm.py

@@ -0,0 +1,250 @@
+import enum
+from typing import Optional
+
+import ixformer.functions as ixff
+import torch
+from ixformer.inference.overlap.linear_mlp_overlap_comm import (
+    LinearMLPOverlapComm,
+    LinearMLPOverlapCommHook,
+)
+
+
+def get_overlap_linear_method(layer):
+    if hasattr(layer, "_overlap_comm_gemm_fn"):
+        return layer._overlap_comm_gemm_fn
+
+    if layer.linear_weights["weight"].itemsize == 2:
+        layer._overlap_comm_gemm_fn = None
+        return None
+
+    def overlap_linear_fn(input, weight, bias=None, out: torch.Tensor = None, **kwargs):
+        return layer.linear_method.apply_weights(
+            layer.linear_weights, input, output=out
+        )
+
+    layer._overlap_comm_gemm_fn = overlap_linear_fn
+    return overlap_linear_fn
+
+
+class DecoderLayerOverlapComm(LinearMLPOverlapCommHook):
+    class HookStage(enum.IntEnum):
+        kExited = 0
+        kTracing = 1
+
+    class HookState:
+        def __init__(self, max_num_chunks):
+            self.max_num_chunks = max_num_chunks
+            self.stage = DecoderLayerOverlapComm.HookStage.kExited
+
+            self.mlp_linaer2_end_events = [
+                torch.cuda.Event() for _ in range(max_num_chunks)
+            ]
+            self.ln_attn_end_event = torch.cuda.Event()
+
+            self.overlap_comm: Optional[LinearMLPOverlapComm] = None
+
+        def is_tracing_stage(self):
+            return self.stage == DecoderLayerOverlapComm.HookStage.kTracing
+
+        def enter(self, overlap_comm, chunk_idx):
+            self.stage = DecoderLayerOverlapComm.HookStage.kTracing
+
+            self.overlap_comm = overlap_comm
+            self.mlp_linaer2_end_events[chunk_idx].record(overlap_comm._comm_stream)
+
+        def exit(self):
+            self.overlap_comm = None
+            self.stage = DecoderLayerOverlapComm.HookStage.kExited
+
+        def __str__(self):
+            return f"HookState(overlap_comm={self.overlap_comm}, stage={self.stage})"
+
+        def __repr__(self):
+            return self.__str__()
+
+    _overlap_comm_hook_state = dict()
+
+    def __init__(self, model_id, layer_idx, max_num_chunks: int = 4):
+        """
+        DecoderLayer 的流程：
+        ln_qkv: InputLayerNorm(hidden_states, [residual]) -> qkv_proj(hidden_states) -> q, k, v = split(hidden_states) -> Attention(q, k, v)
+        linear_mlp: AttentionOutputProj(hidden_states) -> PostLayerNorm(hidden_states) -> MLPLinear1 -> MLPActivation -> MLPLinear2
+
+        其中：AttentionOutputProj 和 MLPLinear2 之后如果使用 TP，那么需要进行 AllReduce
+
+        通过上述流程，该类的目的是将 MLPLinear2 后的 AllReduce 和 DecoderLayer 最开始的 ln_qkv 进行 Overlap。
+        其中，第一层 DecoderLayer 不进行 ln_qkv 的 Overlap，因为在第一层之前没有通讯。
+        我们需要将第 i 层 MLPLinear2 后的通讯 和 第 i + 1 层的 ln_qkv 进行 Overlap。
+
+        为了管理当前的状态和获取前一层的状态，从而设计了 DecoderLayerOverlapComm 类。
+        该类需要 model_id 来推断当前正在运行的模型，用 layer_idx 来标记每一层的开始和结束，
+        以及通过 layer_idx 去获取前一层的状态。
+
+        注：
+            - 在 call_ln_qkv_overlap 中对 Tensor 进行切分时，
+              需要保持和 linear_mlp 切分的大小是一致的，否则会出现 Tensor 的数据不对应；
+            - 如果需要使用 ln_qkv 进行 Overlap，那么必须使用该类的 linear_mlp 去替换 linear_mlp_overlap
+
+        :param model_id: 模型的 id，可以使用 id(model) 去设置
+        :param layer_idx: layer 的索引，注意，需要从 0 到 NumLayers 的顺序去完成构造
+        :param max_num_chunks: 最大能进行切分的次数
+        """
+        self._model_id = model_id
+        self._layer_idx = layer_idx
+        self._max_num_chunks = max_num_chunks
+
+        self._state = self.HookState(max_num_chunks)
+        self._overlap_comm_hook_state[(model_id, layer_idx)] = self._state
+        self._prev_layer_state = (
+            None
+            if layer_idx == 0
+            else self._overlap_comm_hook_state[(model_id, layer_idx - 1)]
+        )
+
+    @property
+    def model_id(self):
+        return self._model_id
+
+    @property
+    def layer_idx(self):
+        return self._layer_idx
+
+    @property
+    def max_num_chunks(self):
+        return self._max_num_chunks
+
+    @property
+    def state(self) -> "DecoderLayerOverlapComm.HookState":
+        return self._state
+
+    @property
+    def prev_layer_state(self) -> "DecoderLayerOverlapComm.HookState":
+        return self._prev_layer_state
+
+    def is_ln_qkv_overlap(self):
+        return not (
+            self.layer_idx == 0
+            or not self.prev_layer_state.is_tracing_stage()
+            or self.prev_layer_state.overlap_comm is None
+        )
+
+    def ln_qkv(self, hidden_states, residual, ln_layer, qkv_layer, out_last_dim):
+        """
+        :param hidden_state: shape[Batch * SeqLen, HiddenSize]
+        :param residual: shape[Batch * SeqLen, HiddenSize]
+        :param ln_layer: torch.nn.Module or Function(hidden_state, residual=None)
+        :param qkv_layer: vllm.QKVParallelLinear
+        :param out_last_dim: qkv_layer 输出 Tensor 的最后一个维度
+        :return: qkv, residual
+        """
+        if self.is_ln_qkv_overlap():
+            qkv, residual = self.call_ln_qkv_overlap(
+                hidden_states, residual, ln_layer, qkv_layer, out_last_dim
+            )
+        else:
+            qkv, residual = self.call_ln_qkv(
+                hidden_states, residual, ln_layer, qkv_layer
+            )
+
+        return qkv, residual
+
+    def call_ln_qkv_overlap(
+        self, hidden_states, residual, ln_layer, qkv_layer, out_last_dim
+    ):
+        if hidden_states.ndim != 2:
+            raise RuntimeError(
+                f"Expected 2-dim for hidden state, but got {hidden_states.ndim}."
+            )
+
+        num_chunks = self.prev_layer_state.overlap_comm.num_chunks
+        overlap_comm: LinearMLPOverlapComm = self.prev_layer_state.overlap_comm
+
+        if num_chunks > self.max_num_chunks:
+            raise RuntimeError(
+                f"The layer is not support more than {self.max_num_chunks}, got {num_chunks}."
+            )
+
+        hidden_state_chunks = list(torch.chunk(hidden_states, num_chunks, dim=0))
+        if residual is None:
+            residual = hidden_states
+            residual_chunks = [None] * num_chunks
+        else:
+            residual_chunks = torch.chunk(residual, num_chunks, dim=0)
+
+        out = torch.empty(
+            (hidden_states.shape[0], out_last_dim),
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+        out_chunks = list(torch.chunk(out, num_chunks, dim=0))
+
+        for chunk_idx, (hidden_state_chunk, residual_chunk, out_chunk) in enumerate(
+            zip(hidden_state_chunks, residual_chunks, out_chunks)
+        ):
+            overlap_comm._compute_streams[
+                chunk_idx % overlap_comm.num_compute_streams
+            ].wait_event(self.prev_layer_state.mlp_linaer2_end_events[chunk_idx])
+            with overlap_comm.compute_stream_context(chunk_idx):
+                self.call_ln_qkv(
+                    hidden_state_chunk,
+                    residual_chunk,
+                    ln_layer,
+                    qkv_layer,
+                    chunk_idx,
+                    use_limited_gemm=chunk_idx != (num_chunks - 1),
+                    out=out_chunk,
+                    overlap_comm=overlap_comm,
+                )
+
+        self.prev_layer_state.exit()
+        overlap_comm.stop_overlap()
+        return out, residual
+
+    def call_ln_qkv(
+        self,
+        hidden_state,
+        residual,
+        ln_layer,
+        qkv_layer,
+        chunk_idx=0,
+        use_limited_gemm=False,
+        out=None,
+        overlap_comm: LinearMLPOverlapComm = None,
+    ):
+        if residual is None:
+            residual = hidden_state
+            if ln_layer is not None:
+                hidden_state = ln_layer(hidden_state)
+        else:
+            hidden_state, residual = ln_layer(hidden_state, residual)
+
+        if out is None:
+            qkv, _ = qkv_layer(hidden_state)
+        else:
+            gemm_method = get_overlap_linear_method(qkv_layer)
+            qkv = overlap_comm.gemm_dispatcher(
+                chunk_idx=chunk_idx,
+                chunk_input=hidden_state,
+                weight=qkv_layer.linear_weights["weight"],
+                chunk_out=out,
+                user_gemm_method=gemm_method,
+                use_limited_gemm=use_limited_gemm,
+            )
+
+        return qkv, residual
+
+    def linear_mlp(self, *args, **kwargs):
+        """ref: linear_mlp_overlap"""
+        return ixff.linear_mlp_overlap(
+            *args, **kwargs, mlp_linear2_finished_callback=self.on_mlp_linear2_finished
+        )
+
+    def on_mlp_linear2_finished(
+        self,
+        overlap_comm: LinearMLPOverlapComm,
+        num_chunks,
+        chunk_idx,
+        hidden_states_chunk,
+        residual_chunk,
+    ):
+        self.state.enter(overlap_comm, chunk_idx)

ixformer_sdk/inference/overlap/w8a8_allreduce.py

@@ -0,0 +1,154 @@
+import math
+from typing import Optional
+
+import ixformer.distributed as ixfd
+import ixformer.functions as F
+import torch
+import torch.distributed as dist
+from ixformer.distributed.overlap_comm import SplitOverlapComm
+
+from ixformer.core import config as ixff_config
+
+__all__ = ["w8a8_allreduce"]
+
+
+class W8A8AllReduceOverlap(SplitOverlapComm):
+    def compute(
+        self,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        input_scale: torch.Tensor,
+        weight_scale: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        output: Optional[torch.Tensor] = None,
+        format: str = "TN",
+        out_dtype: torch.dtype = None,
+        comm_group=None,
+        split_ratio=0.5,
+    ):
+        # compute the chunk size of input
+        input_chunk_sizes = [int(math.ceil(input.shape[0] * split_ratio))]
+        input_chunk_sizes.append(input.shape[0] - input_chunk_sizes[0])
+
+        # split input and input_scale
+        input_chunks = list(torch.split_with_sizes(input, input_chunk_sizes, dim=0))
+        input_scale_chunks = torch.split(
+            input_scale,
+            input_chunk_sizes,
+        )
+
+        # create output and split it
+        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,
+            )
+        out_chunks = torch.split(output, input_chunk_sizes)
+
+        # overlap gemm and allreduce
+        for chunk_idx in range(len(input_chunks)):
+            # submit gemm kernel into compute stream
+            with self.compute_stream_context(chunk_idx):
+                F.w8a8(
+                    input=input_chunks[chunk_idx],
+                    weight=weight,
+                    i_scales=input_scale_chunks[chunk_idx],
+                    w_scales=weight_scale,
+                    bias=bias,
+                    output=out_chunks[chunk_idx],
+                    format=format,
+                    persistent=chunk_idx != 0,
+                )
+
+            # recode compute stream and wait gemm
+            self.start_comm(chunk_idx)
+
+            # submit allreduce kernel into communication stream by set use_comm_stream to true
+            ixfd.all_reduce(
+                out_chunks[chunk_idx],
+                async_op=True,
+                group=self.comm_group,
+                use_comm_stream=True,
+            )
+
+        return output
+
+
+_w8a8_allreduce_overlap = None
+
+
+def w8a8_allreduce(
+    enable_overlap: bool,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+    output: Optional[torch.Tensor] = None,
+    format: str = "TN",
+    out_dtype: torch.dtype = None,
+    comm_group=None,
+    split_ratio=0.5,
+) -> torch.Tensor:
+    """
+    Gemm(w8a8) + AllReduce
+
+    Args:
+        enable_overlap: whether enable gemm and allreduce overlap
+        input:          shape: [M, K], dtype: int8,                linear input
+        weight:         shape: [N, K], dtype: int8,                linear weight
+        input_scale:    shape: [M],    dtype: float32,             quantized scale of input
+        weight_scale:   shape: [N],    dtype: float32,             quantized scale of weight
+        bias:           shape: [N],    dtype: float16 or bfloat16, linear bias
+        output:         shape: [M, N], dtype: float16 or bfloat16, allreduce output
+        format:         options include TN, NN and NT
+        out_dtype:      use the argument to decide to the dtype of output when output is None
+        comm_group:     communication group
+        split_ratio:    split the ratio of input.shape[0] when using overlap, range: (0, 1),
+                        it will affect area of the overlap for gemm and allreduce.
+    Returns: output
+    """
+    if (
+        enable_overlap
+        and ixff_config.IXFORMER_ENABLE_OVERLAP_COMM
+        and dist.is_initialized()
+        and dist.get_world_size(comm_group) > 1
+        and input.shape[0] > 1
+    ):
+        global _w8a8_allreduce_overlap
+        if _w8a8_allreduce_overlap is None:
+            _w8a8_allreduce_overlap = W8A8AllReduceOverlap.dispatcher(
+                num_chunks=2, comm_group=comm_group
+            ).forward
+        return _w8a8_allreduce_overlap(
+            input=input,
+            weight=weight,
+            input_scale=input_scale,
+            weight_scale=weight_scale,
+            bias=bias,
+            output=output,
+            format=format,
+            out_dtype=out_dtype,
+            split_ratio=split_ratio,
+        )
+
+    out = F.w8a8(
+        input=input,
+        weight=weight,
+        i_scales=input_scale,
+        w_scales=weight_scale,
+        bias=bias,
+        output=output,
+        format=format,
+        out_dtype=out_dtype,
+    )
+
+    if dist.get_world_size() > 1:
+        ixfd.all_reduce(out, op=ixfd.ReduceOp.SUM, async_op=True, group=comm_group)
+
+    return out