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/.gitignore

@@ -0,0 +1,2 @@
+*.so
+build/

ixformer_sdk/__init__.py

@@ -0,0 +1,2 @@
+import torch
+from .functions import *

ixformer_sdk/contrib/DeepCache/__init__.py

@@ -0,0 +1,5 @@
+from .sd.pipeline_stable_diffusion import StableDiffusionPipeline
+from .sdxl.pipeline_stable_diffusion_xl import StableDiffusionXLPipeline
+from .sdxl.pipeline_stable_diffusion_xl_img2img import StableDiffusionXLImg2ImgPipeline
+
+from .sd.pipeline_text_to_video_zero import TextToVideoZeroPipeline

ixformer_sdk/contrib/DeepCache/ddpm/ddim.py

@@ -0,0 +1,180 @@
+import argparse
+import traceback
+import shutil
+import logging
+import yaml
+import random
+import sys
+import os
+import torch
+import numpy as np
+
+from ddpm.utils.logging import Logger, EmptyLogger
+from ddpm.utils.tools import set_random_seed
+from accelerate import Accelerator, DistributedDataParallelKwargs
+
+torch.set_printoptions(sci_mode=False)
+
+def dict2namespace(config):
+    namespace = argparse.Namespace()
+    for key, value in config.items():
+        if isinstance(value, dict):
+            new_value = dict2namespace(value)
+        else:
+            new_value = value
+        setattr(namespace, key, new_value)
+    return namespace
+
+def parse_args_and_config():
+    parser = argparse.ArgumentParser(description=globals()["__doc__"])
+
+    parser.add_argument(
+        "--config", type=str, required=True, help="Path to the config file"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=1234, help="Random seed")
+    parser.add_argument(
+        "--exp", type=str, default="exp", help="Path for saving running related data."
+    )
+    
+    parser.add_argument(
+        "--test", action="store_true", help="Whether to test the model"
+    )
+    parser.add_argument(
+        "--sample", action="store_true", help="Whether to produce samples from the model",
+    )
+    parser.add_argument(
+        "--image_folder", type=str, default="images", help="folder name for storing the sampled images"
+    )
+
+    parser.add_argument(
+        "--fid", action="store_true"
+    )
+    parser.add_argument(
+        "--interpolation", action="store_true"
+    )
+    parser.add_argument(
+        "--resume_training", action="store_true", help="Whether to resume training"
+    )
+    parser.add_argument(
+        "--ni", action="store_true", help="No interaction. Suitable for Slurm Job launcher",
+    )
+    parser.add_argument(
+        "--use_pretrained", action="store_true"
+    )
+    parser.add_argument(
+        "--sample_type", type=str, default="generalized", help="sampling approach (generalized or ddpm_noisy)",
+    )
+    parser.add_argument(
+        "--skip_type", type=str, default="uniform", help="skip according to (uniform or quadratic)",
+    )
+    parser.add_argument(
+        "--timesteps", type=int, default=1000, help="number of steps involved"
+    )
+    parser.add_argument(
+        "--eta", type=float, default=0.0, help="eta used to control the variances of sigma",
+    )
+    parser.add_argument(
+        "--dyn", action="store_true", help="whether to activate the dynamic train/inference"
+    )
+    parser.add_argument(
+        "--sequence", action="store_true"
+    )
+    parser.add_argument(
+        "--select_step", type=int, default=None
+    )
+    parser.add_argument(
+        "--select_depth", type=int, default=None
+    )
+
+    parser.add_argument(
+        "--cache", action="store_true"
+    )
+    parser.add_argument(
+        "--cache_interval", type=int, default=None,
+    )
+    parser.add_argument(
+        "--non_uniform", action="store_true"
+    )
+    parser.add_argument(
+        "--pow", type=float, default=None,
+    )
+    parser.add_argument(
+        "--center", type=int, default=None,
+    )
+    parser.add_argument(
+        "--branch", type=int, default=None,
+    )
+    args = parser.parse_args()
+    # parse config file
+    with open(args.config, "r") as f:
+        config = yaml.safe_load(f)
+    new_config = dict2namespace(config)
+    new_config.select_step = args.select_step
+    new_config.select_depth = args.select_depth
+
+    torch.backends.cudnn.benchmark = True
+
+    return args, new_config
+
+
+def main():
+    args, config = parse_args_and_config()
+    
+    if args.dyn:
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+        accelerator = Accelerator(kwargs_handlers=[ddp_kwargs])
+    else:
+        accelerator = Accelerator()
+    args.accelerator = accelerator
+
+    #log_root_dir = "{}_runtime_log".format(args.config[8:-4])
+    log_root_dir = "runtime_log"
+    dataset = args.config[8:-4]
+    if args.cache:
+        if args.non_uniform:
+            sub_dir_name = "{}_{}_cache_{}_pow_{}_center_{}".format(dataset, args.exp, args.cache_interval, args.pow, args.center)
+        else:
+            sub_dir_name = "{}_{}_cache_{}".format(dataset, args.exp, args.cache_interval)
+    else:
+        sub_dir_name = "{}".format(args.exp)
+
+    if accelerator.is_main_process:
+        logger = Logger(
+            root_dir=log_root_dir,
+            sub_name=sub_dir_name,
+            config=args.__dict__, 
+            append=(args.sample == True)
+        )
+        args.logger = logger
+
+        args.logger.log("Writing log file to {}".format(args.logger.sub_dir))
+        args.logger.log("Exp instance PID = {}".format(os.getpid()))
+    else:
+        args.logger = EmptyLogger(
+            root_dir=log_root_dir,
+            sub_name=sub_dir_name,
+        )
+
+    args.image_folder = args.logger.setup_image_folder("{}".format(args.image_folder))
+
+    args.seed += accelerator.process_index      
+    # set random seed
+    set_random_seed(args.seed)
+    try:
+        if args.cache:
+            from ddpm.runners.deepcache import Diffusion
+            runner = Diffusion(args, config)
+            runner.sample()
+        else:
+            from ddpm.runners.diffusion import Diffusion
+            runner = Diffusion(args, config)
+            runner.sample()
+    except Exception:
+        logging.error(traceback.format_exc())
+
+    return 0
+
+
+if __name__ == "__main__":
+    main()

ixformer_sdk/contrib/DeepCache/ddpm/fid.py

@@ -0,0 +1,361 @@
+"""Calculates the Frechet Inception Distance (FID) to evalulate GANs
+
+The FID metric calculates the distance between two distributions of images.
+Typically, we have summary statistics (mean & covariance matrix) of one
+of these distributions, while the 2nd distribution is given by a GAN.
+
+When run as a stand-alone program, it compares the distribution of
+images that are stored as PNG/JPEG at a specified location with a
+distribution given by summary statistics (in pickle format).
+
+The FID is calculated by assuming that X_1 and X_2 are the activations of
+the pool_3 layer of the inception net for generated samples and real world
+samples respectively.
+
+See --help to see further details.
+
+Code apapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead
+of Tensorflow
+
+Copyright 2018 Institute of Bioinformatics, JKU Linz
+
+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.
+"""
+import os
+import pathlib
+from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
+
+import numpy as np
+import torch
+import torchvision.transforms as TF
+from PIL import Image
+from scipy import linalg
+from torch.nn.functional import adaptive_avg_pool2d
+
+try:
+    from tqdm import tqdm
+except ImportError:
+    # If tqdm is not available, provide a mock version of it
+    def tqdm(x):
+        return x
+
+from pytorch_fid.inception import InceptionV3
+
+parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
+parser.add_argument('--batch-size', type=int, default=50,
+                    help='Batch size to use')
+parser.add_argument('--dataset_name', type=str, default=None)
+parser.add_argument('--num-workers', type=int,
+                    help=('Number of processes to use for data loading. '
+                          'Defaults to `min(8, num_cpus)`'))
+parser.add_argument('--device', type=str, default=None,
+                    help='Device to use. Like cuda, cuda:0 or cpu')
+parser.add_argument('--dims', type=int, default=2048,
+                    choices=list(InceptionV3.BLOCK_INDEX_BY_DIM),
+                    help=('Dimensionality of Inception features to use. '
+                          'By default, uses pool3 features'))
+parser.add_argument('--num_samples', type=int, default=None,
+                    help=('Number of samples for FID estimation'))
+parser.add_argument('--res', type=int, default=None,
+                    help=('Resolutions of samples for FID estimation'))
+parser.add_argument('--save-stats', action='store_true',
+                    help=('Generate an npz archive from a directory of samples. '
+                          'The first path is used as input and the second as output.'))
+
+parser.add_argument('--path', type=str, nargs=2,
+                    help=('Paths to the generated images or '
+                          'to .npz statistic files'))
+
+
+IMAGE_EXTENSIONS = {'bmp', 'jpg', 'jpeg', 'pgm', 'png', 'ppm',
+                    'tif', 'tiff', 'webp'}
+
+
+class ImagePathDataset(torch.utils.data.Dataset):
+    def __init__(self, files, transforms=None):
+        self.files = files
+        self.transforms = transforms
+
+    def __len__(self):
+        return len(self.files)
+
+    def __getitem__(self, i):
+        path = self.files[i]
+        img = Image.open(path).convert('RGB')
+        if self.transforms is not None:
+            img = self.transforms(img)
+        return img
+
+
+def get_activations(files, model, batch_size=50, dims=2048, device='cpu',
+                    num_workers=1, res=None, dataset_name=None):
+    """Calculates the activations of the pool_3 layer for all images.
+
+    Params:
+    -- files       : List of image files paths
+    -- model       : Instance of inception model
+    -- batch_size  : Batch size of images for the model to process at once.
+                     Make sure that the number of samples is a multiple of
+                     the batch size, otherwise some samples are ignored. This
+                     behavior is retained to match the original FID score
+                     implementation.
+    -- dims        : Dimensionality of features returned by Inception
+    -- device      : Device to run calculations
+    -- num_workers : Number of parallel dataloader workers
+
+    Returns:
+    -- A numpy array of dimension (num images, dims) that contains the
+       activations of the given tensor when feeding inception with the
+       query tensor.
+    """
+    model.eval()
+
+    if batch_size > len(files):
+        print(('Warning: batch size is bigger than the data size. '
+               'Setting batch size to data size'))
+        batch_size = len(files)
+
+    if res is None:
+        trans = TF.ToTensor()
+    else:
+        if dataset_name == 'celeba':
+            from switchable_diffusion.datasets import Crop
+            print("In crop image: {}, {}".format(res, dataset_name))
+            cx = 89
+            cy = 121
+            x1 = cy - 64
+            x2 = cy + 64
+            y1 = cx - 64
+            y2 = cx + 64
+            trans = TF.Compose([
+                        Crop(x1, x2, y1, y2),
+                        TF.Resize(res),
+                        TF.ToTensor(),
+            ])
+        else:
+            trans = TF.Compose([
+                TF.Resize(res),
+                TF.CenterCrop(res),
+                TF.ToTensor()
+            ])
+    
+    dataset = ImagePathDataset(files, transforms=trans)
+    dataloader = torch.utils.data.DataLoader(dataset,
+                                             batch_size=batch_size,
+                                             shuffle=False,
+                                             drop_last=False,
+                                             num_workers=num_workers)
+
+    pred_arr = np.empty((len(files), dims))
+
+    start_idx = 0
+
+    for batch in tqdm(dataloader):
+        batch = batch.to(device)
+
+        with torch.no_grad():
+            pred = model(batch)[0]
+
+        # If model output is not scalar, apply global spatial average pooling.
+        # This happens if you choose a dimensionality not equal 2048.
+        if pred.size(2) != 1 or pred.size(3) != 1:
+            pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
+
+        pred = pred.squeeze(3).squeeze(2).cpu().numpy()
+
+        pred_arr[start_idx:start_idx + pred.shape[0]] = pred
+
+        start_idx = start_idx + pred.shape[0]
+
+    return pred_arr
+
+
+def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+    """Numpy implementation of the Frechet Distance.
+    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+    and X_2 ~ N(mu_2, C_2) is
+            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+
+    Stable version by Dougal J. Sutherland.
+
+    Params:
+    -- mu1   : Numpy array containing the activations of a layer of the
+               inception net (like returned by the function 'get_predictions')
+               for generated samples.
+    -- mu2   : The sample mean over activations, precalculated on an
+               representative data set.
+    -- sigma1: The covariance matrix over activations for generated samples.
+    -- sigma2: The covariance matrix over activations, precalculated on an
+               representative data set.
+
+    Returns:
+    --   : The Frechet Distance.
+    """
+
+    mu1 = np.atleast_1d(mu1)
+    mu2 = np.atleast_1d(mu2)
+
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+
+    assert mu1.shape == mu2.shape, \
+        'Training and test mean vectors have different lengths'
+    assert sigma1.shape == sigma2.shape, \
+        'Training and test covariances have different dimensions'
+
+    diff = mu1 - mu2
+
+    # Product might be almost singular
+    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    if not np.isfinite(covmean).all():
+        msg = ('fid calculation produces singular product; '
+               'adding %s to diagonal of cov estimates') % eps
+        print(msg)
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+    # Numerical error might give slight imaginary component
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            m = np.max(np.abs(covmean.imag))
+            raise ValueError('Imaginary component {}'.format(m))
+        covmean = covmean.real
+
+    tr_covmean = np.trace(covmean)
+
+    return (diff.dot(diff) + np.trace(sigma1)
+            + np.trace(sigma2) - 2 * tr_covmean)
+
+
+def calculate_activation_statistics(files, model, batch_size=50, dims=2048,
+                                    device='cpu', num_workers=1, res=None, dataset_name=None):
+    """Calculation of the statistics used by the FID.
+    Params:
+    -- files       : List of image files paths
+    -- model       : Instance of inception model
+    -- batch_size  : The images numpy array is split into batches with
+                     batch size batch_size. A reasonable batch size
+                     depends on the hardware.
+    -- dims        : Dimensionality of features returned by Inception
+    -- device      : Device to run calculations
+    -- num_workers : Number of parallel dataloader workers
+
+    Returns:
+    -- mu    : The mean over samples of the activations of the pool_3 layer of
+               the inception model.
+    -- sigma : The covariance matrix of the activations of the pool_3 layer of
+               the inception model.
+    """
+    act = get_activations(files, model, batch_size, dims, device, num_workers, res=res, dataset_name=dataset_name)
+    mu = np.mean(act, axis=0)
+    sigma = np.cov(act, rowvar=False)
+    return mu, sigma
+
+
+def compute_statistics_of_path(path, model, batch_size, dims, device,
+                               num_workers=1, num_samples=None, res=None, dataset_name=None):
+    if path.endswith('.npz'):
+        with np.load(path) as f:
+            m, s = f['mu'][:], f['sigma'][:]
+    else:
+        path = pathlib.Path(path)
+
+        files = sorted([file for ext in IMAGE_EXTENSIONS
+                       for file in path.glob('**/*.{}'.format(ext))])
+        if num_samples is not None:
+            #import random
+            #files = random.sample(files, num_samples)
+            files = files[:num_samples]
+        print("Found %d files." % len(files))
+        m, s = calculate_activation_statistics(files, model, batch_size,
+                                               dims, device, num_workers, res=res, dataset_name=dataset_name)
+
+    return m, s
+
+
+def calculate_fid_given_paths(paths, batch_size, device, dims, num_workers=1, num_samples=None, res=None, dataset_name=None):
+    """Calculates the FID of two paths"""
+    for p in paths:
+        if not os.path.exists(p):
+            raise RuntimeError('Invalid path: %s' % p)
+
+    block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
+
+    model = InceptionV3([block_idx]).to(device)
+
+    m1, s1 = compute_statistics_of_path(paths[0], model, batch_size,
+                                        dims, device, num_workers, num_samples=num_samples, res=res, dataset_name=dataset_name)
+    m2, s2 = compute_statistics_of_path(paths[1], model, batch_size,
+                                        dims, device, num_workers, num_samples=num_samples, res=res, dataset_name=dataset_name)
+    fid_value = calculate_frechet_distance(m1, s1, m2, s2)
+
+    return fid_value
+
+
+def save_fid_stats(paths, batch_size, device, dims, num_workers=1, num_samples=None, res=None, dataset_name=None):
+    """Calculates the FID of two paths"""
+    if not os.path.exists(paths[0]):
+        raise RuntimeError('Invalid path: %s' % paths[0])
+
+    if os.path.exists(paths[1]):
+        raise RuntimeError('Existing output file: %s' % paths[1])
+
+    block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
+
+    model = InceptionV3([block_idx]).to(device)
+
+    print(f"Saving statistics for {paths[0]}")
+
+    m1, s1 = compute_statistics_of_path(paths[0], model, batch_size,
+                                        dims, device, num_workers, num_samples=num_samples, res=res, dataset_name=dataset_name)
+
+    np.savez_compressed(paths[1], mu=m1, sigma=s1)
+
+
+def main():
+    args = parser.parse_args()
+
+    if args.device is None:
+        device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')
+    else:
+        device = torch.device(args.device)
+
+    if args.num_workers is None:
+        try:
+            num_cpus = len(os.sched_getaffinity(0))
+        except AttributeError:
+            # os.sched_getaffinity is not available under Windows, use
+            # os.cpu_count instead (which may not return the *available* number
+            # of CPUs).
+            num_cpus = os.cpu_count()
+
+        num_workers = min(num_cpus, 8) if num_cpus is not None else 0
+    else:
+        num_workers = args.num_workers
+
+    if args.save_stats:
+        save_fid_stats(args.path, args.batch_size, device, args.dims, num_workers, num_samples=args.num_samples, res=args.res, dataset_name=args.dataset_name)
+        return
+
+    fid_value = calculate_fid_given_paths(args.path,
+                                          args.batch_size,
+                                          device,
+                                          args.dims,
+                                          num_workers,
+                                          num_samples=args.num_samples,
+                                          res = args.res, dataset_name=args.dataset_name)
+    print('FID: ', fid_value)
+
+
+if __name__ == '__main__':
+    main()

ixformer_sdk/contrib/DeepCache/flops.py

@@ -0,0 +1,559 @@
+'''
+This opcounter is adapted from https://github.com/sovrasov/flops-counter.pytorch and https://github.com/Lyken17/pytorch-OpCounter
+
+Copyright (C) 2021 Sovrasov V. - All Rights Reserved
+ * You may use, distribute and modify this code under the
+ * terms of the MIT license.
+ * You should have received a copy of the MIT license with
+ * this file. If not visit https://opensource.org/licenses/MIT
+'''
+import os
+import yaml
+import numpy as np
+import torch.nn as nn
+import torch
+has_timm = False
+
+from diffusers.models.lora import LoRACompatibleLinear, LoRACompatibleConv
+
+@torch.no_grad()
+def count_ops_and_params(model, example_inputs, layer_wise=False):
+    global CUSTOM_MODULES_MAPPING
+    ori_model = model 
+    model = copy.deepcopy(model) # deepcopy to avoid changing the original model
+    flops_model = add_flops_counting_methods(model)
+    flops_model.eval()
+    flops_model.start_flops_count(ost=sys.stdout, verbose=False,
+                                ignore_list=[])
+    if isinstance(example_inputs, (tuple, list)):
+        _ = flops_model(*example_inputs)
+    elif isinstance(example_inputs, dict):
+        _ = flops_model(**example_inputs)
+    else:
+        _ = flops_model(example_inputs)
+    flops_count, params_count, _layer_flops, _layer_params = flops_model.compute_average_flops_cost()
+    layer_flops = {}
+    layer_params = {}
+
+    for m_name, m in model.named_modules():
+        layer_flops[m_name] = _layer_flops.get(m)
+        layer_params[m_name] = _layer_params.get(m)
+        if layer_wise:
+            space = 30 - len(m_name)
+            print("Layer {}: {} MACs = {:.4f} G, Params = {:.4f} M, MACs% = {:.2f}".format(
+                m_name, ' ' * space, layer_flops[m_name]/1e9, layer_params[m_name] / 1e6, 100 * layer_flops[m_name] / flops_count
+            ))
+
+    flops_model.stop_flops_count()
+    CUSTOM_MODULES_MAPPING = {}
+    #if layer_wise:
+    #    return flops_count, params_count, layer_flops, layer_params
+    return flops_count, params_count
+
+def empty_flops_counter_hook(module, input, output):
+    module.__flops__ += 0
+
+
+def upsample_flops_counter_hook(module, input, output):
+    output_size = output[0]
+    batch_size = output_size.shape[0]
+    output_elements_count = batch_size
+    for val in output_size.shape[1:]:
+        output_elements_count *= val
+    module.__flops__ += int(output_elements_count)
+
+
+def relu_flops_counter_hook(module, input, output):
+    active_elements_count = output.numel()
+    module.__flops__ += int(active_elements_count)
+
+
+def linear_flops_counter_hook(module, input, output):
+    input = input[0]
+    # pytorch checks dimensions, so here we don't care much
+    output_last_dim = output.shape[-1]
+    bias_flops = output_last_dim if module.bias is not None else 0
+    module.__flops__ += int(np.prod(input.shape) * output_last_dim + bias_flops)
+
+
+def pool_flops_counter_hook(module, input, output):
+    input = input[0]
+    module.__flops__ += int(np.prod(input.shape))
+
+
+def bn_flops_counter_hook(module, input, output):
+    input = input[0]
+
+    batch_flops = np.prod(input.shape)
+    if module.affine:
+        batch_flops *= 2
+    module.__flops__ += int(batch_flops)
+
+def ln_flops_counter_hook(module, input, output):
+    input = input[0]
+    batch_flops = np.prod(input.shape)
+    if module.elementwise_affine:
+        batch_flops *= 2
+    module.__flops__ += int(batch_flops)
+
+def conv_flops_counter_hook(conv_module, input, output):
+    # Can have multiple inputs, getting the first one
+    input = input[0]
+
+    batch_size = input.shape[0]
+    output_dims = list(output.shape[2:])
+
+    kernel_dims = list(conv_module.kernel_size)
+    in_channels = conv_module.in_channels
+    out_channels = conv_module.out_channels
+    groups = conv_module.groups
+
+    filters_per_channel = out_channels // groups
+    conv_per_position_flops = int(np.prod(kernel_dims)) * \
+        in_channels * filters_per_channel
+
+    active_elements_count = batch_size * int(np.prod(output_dims))
+
+    overall_conv_flops = conv_per_position_flops * active_elements_count
+
+    bias_flops = 0
+
+    if conv_module.bias is not None:
+
+        bias_flops = out_channels * active_elements_count
+
+    overall_flops = overall_conv_flops + bias_flops
+
+    conv_module.__flops__ += int(overall_flops)
+
+
+def rnn_flops(flops, rnn_module, w_ih, w_hh, input_size):
+    # matrix matrix mult ih state and internal state
+    flops += w_ih.shape[0]*w_ih.shape[1]
+    # matrix matrix mult hh state and internal state
+    flops += w_hh.shape[0]*w_hh.shape[1]
+    if isinstance(rnn_module, (nn.RNN, nn.RNNCell)):
+        # add both operations
+        flops += rnn_module.hidden_size
+    elif isinstance(rnn_module, (nn.GRU, nn.GRUCell)):
+        # hadamard of r
+        flops += rnn_module.hidden_size
+        # adding operations from both states
+        flops += rnn_module.hidden_size*3
+        # last two hadamard product and add
+        flops += rnn_module.hidden_size*3
+    elif isinstance(rnn_module, (nn.LSTM, nn.LSTMCell)):
+        # adding operations from both states
+        flops += rnn_module.hidden_size*4
+        # two hadamard product and add for C state
+        flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size
+        # final hadamard
+        flops += rnn_module.hidden_size + rnn_module.hidden_size + rnn_module.hidden_size
+    return flops
+
+
+def rnn_flops_counter_hook(rnn_module, input, output):
+    """
+    Takes into account batch goes at first position, contrary
+    to pytorch common rule (but actually it doesn't matter).
+    If sigmoid and tanh are hard, only a comparison FLOPS should be accurate
+    """
+    flops = 0
+    # input is a tuple containing a sequence to process and (optionally) hidden state
+    inp = input[0]
+    batch_size = inp[0].shape[0]
+    seq_length = inp[0].shape[1]
+    num_layers = rnn_module.num_layers
+
+    for i in range(num_layers):
+        w_ih = rnn_module.__getattr__('weight_ih_l' + str(i))
+        w_hh = rnn_module.__getattr__('weight_hh_l' + str(i))
+        if i == 0:
+            input_size = rnn_module.input_size
+        else:
+            input_size = rnn_module.hidden_size
+        flops = rnn_flops(flops, rnn_module, w_ih, w_hh, input_size)
+        if rnn_module.bias:
+            b_ih = rnn_module.__getattr__('bias_ih_l' + str(i))
+            b_hh = rnn_module.__getattr__('bias_hh_l' + str(i))
+            flops += b_ih.shape[0] + b_hh.shape[0]
+
+    flops *= batch_size
+    flops *= seq_length
+    if rnn_module.bidirectional:
+        flops *= 2
+    rnn_module.__flops__ += int(flops)
+
+
+def rnn_cell_flops_counter_hook(rnn_cell_module, input, output):
+    flops = 0
+    inp = input[0]
+    batch_size = inp.shape[0]
+    w_ih = rnn_cell_module.__getattr__('weight_ih')
+    w_hh = rnn_cell_module.__getattr__('weight_hh')
+    input_size = inp.shape[1]
+    flops = rnn_flops(flops, rnn_cell_module, w_ih, w_hh, input_size)
+    if rnn_cell_module.bias:
+        b_ih = rnn_cell_module.__getattr__('bias_ih')
+        b_hh = rnn_cell_module.__getattr__('bias_hh')
+        flops += b_ih.shape[0] + b_hh.shape[0]
+
+    flops *= batch_size
+    rnn_cell_module.__flops__ += int(flops)
+
+
+def multihead_attention_counter_hook(multihead_attention_module, input, output):
+    flops = 0
+    q, k, v = input
+
+    batch_first = multihead_attention_module.batch_first \
+        if hasattr(multihead_attention_module, 'batch_first') else False
+    if batch_first:
+        batch_size = q.shape[0]
+        len_idx = 1
+    else:
+        batch_size = q.shape[1]
+        len_idx = 0
+
+    dim_idx = 2
+
+    qdim = q.shape[dim_idx]
+    kdim = k.shape[dim_idx]
+    vdim = v.shape[dim_idx]
+
+    qlen = q.shape[len_idx]
+    klen = k.shape[len_idx]
+    vlen = v.shape[len_idx]
+
+    num_heads = multihead_attention_module.num_heads
+    assert qdim == multihead_attention_module.embed_dim
+
+    if multihead_attention_module.kdim is None:
+        assert kdim == qdim
+    if multihead_attention_module.vdim is None:
+        assert vdim == qdim
+
+    flops = 0
+
+    # Q scaling
+    flops += qlen * qdim
+    # Initial projections
+    flops += (
+        (qlen * qdim * qdim)  # QW
+        + (klen * kdim * kdim)  # KW
+        + (vlen * vdim * vdim)  # VW
+    )
+    if multihead_attention_module.in_proj_bias is not None:
+        flops += (qlen + klen + vlen) * qdim
+    # attention heads: scale, matmul, softmax, matmul
+    qk_head_dim = qdim // num_heads
+    v_head_dim = vdim // num_heads
+
+    head_flops = (
+        (qlen * klen * qk_head_dim)  # QK^T
+        + (qlen * klen)  # softmax
+        + (qlen * klen * v_head_dim)  # AV
+    )
+    flops += num_heads * head_flops
+    # final projection, bias is always enabled
+    flops += qlen * vdim * (vdim + 1)
+    flops *= batch_size
+    multihead_attention_module.__flops__ += int(flops)
+
+def timm_multihead_attention_counter_hook(multihead_attention_module, input, output):
+    flops = 0
+    
+    q, k, v = input[0], input[0], input[0]
+    input_dim = input[0].shape[2]
+    input_len = input[0].shape[1]
+    batch_size = input[0].shape[0]
+
+    kdim = qdim = vdim = multihead_attention_module.qkv.out_features//3
+    qlen = klen = vlen = input_len
+
+    num_heads = multihead_attention_module.num_heads
+    assert qdim == multihead_attention_module.head_dim * multihead_attention_module.num_heads
+    
+    flops = 0
+    # Q scaling
+    flops += qlen * qdim
+    # Initial projections
+    flops += (
+        (qlen * input_dim * qdim)  # QW
+        + (klen * input_dim * kdim)  # KW
+        + (vlen * input_dim * vdim)  # VW
+    )
+
+    if multihead_attention_module.qkv.bias is not None:
+        flops += (qlen + klen + vlen) * qdim
+    # attention heads: scale, matmul, softmax, matmul
+    qk_head_dim = qdim // num_heads
+    v_head_dim = vdim // num_heads
+
+    head_flops = (
+        (qlen * klen * qk_head_dim)  # QK^T
+        + (qlen * klen)  # softmax
+        + (qlen * klen * v_head_dim)  # AV
+    )
+    flops += num_heads * head_flops
+    # final projection, bias is always enabled
+    flops += qlen * vdim * (vdim + 1)
+    flops *= batch_size
+    multihead_attention_module.__flops__ += int(flops)
+
+
+
+CUSTOM_MODULES_MAPPING = {}
+
+MODULES_MAPPING = {
+    # convolutions
+    nn.Conv1d: conv_flops_counter_hook,
+    nn.Conv2d: conv_flops_counter_hook,
+    nn.Conv3d: conv_flops_counter_hook,
+    LoRACompatibleConv: conv_flops_counter_hook,
+    # activations
+    nn.ReLU: relu_flops_counter_hook,
+    nn.PReLU: relu_flops_counter_hook,
+    nn.ELU: relu_flops_counter_hook,
+    nn.LeakyReLU: relu_flops_counter_hook,
+    nn.ReLU6: relu_flops_counter_hook,
+    # poolings
+    nn.MaxPool1d: pool_flops_counter_hook,
+    nn.AvgPool1d: pool_flops_counter_hook,
+    nn.AvgPool2d: pool_flops_counter_hook,
+    nn.MaxPool2d: pool_flops_counter_hook,
+    nn.MaxPool3d: pool_flops_counter_hook,
+    nn.AvgPool3d: pool_flops_counter_hook,
+    nn.AdaptiveMaxPool1d: pool_flops_counter_hook,
+    nn.AdaptiveAvgPool1d: pool_flops_counter_hook,
+    nn.AdaptiveMaxPool2d: pool_flops_counter_hook,
+    nn.AdaptiveAvgPool2d: pool_flops_counter_hook,
+    nn.AdaptiveMaxPool3d: pool_flops_counter_hook,
+    nn.AdaptiveAvgPool3d: pool_flops_counter_hook,
+    # BNs
+    nn.BatchNorm1d: bn_flops_counter_hook,
+    nn.BatchNorm2d: bn_flops_counter_hook,
+    nn.BatchNorm3d: bn_flops_counter_hook,
+
+    nn.InstanceNorm1d: bn_flops_counter_hook,
+    nn.InstanceNorm2d: bn_flops_counter_hook,
+    nn.InstanceNorm3d: bn_flops_counter_hook,
+    nn.GroupNorm: bn_flops_counter_hook,
+    nn.LayerNorm: ln_flops_counter_hook,
+    # FC
+    nn.Linear: linear_flops_counter_hook,
+    LoRACompatibleLinear: linear_flops_counter_hook,
+    # Upscale
+    nn.Upsample: upsample_flops_counter_hook,
+    # Deconvolution
+    nn.ConvTranspose1d: conv_flops_counter_hook,
+    nn.ConvTranspose2d: conv_flops_counter_hook,
+    nn.ConvTranspose3d: conv_flops_counter_hook,
+    # RNN
+    nn.RNN: rnn_flops_counter_hook,
+    nn.GRU: rnn_flops_counter_hook,
+    nn.LSTM: rnn_flops_counter_hook,
+    nn.RNNCell: rnn_cell_flops_counter_hook,
+    nn.LSTMCell: rnn_cell_flops_counter_hook,
+    nn.GRUCell: rnn_cell_flops_counter_hook,
+    nn.MultiheadAttention: multihead_attention_counter_hook
+}
+
+if has_timm:
+    MODULES_MAPPING.update(
+        {
+            timm.models.vision_transformer.Attention: timm_multihead_attention_counter_hook,
+        }
+    )
+
+if hasattr(nn, 'GELU'):
+    MODULES_MAPPING[nn.GELU] = relu_flops_counter_hook
+
+
+import sys
+from functools import partial
+import torch.nn as nn
+import copy
+
+def accumulate_flops(self, layer_flops):
+    if is_supported_instance(self):
+        layer_flops[self] = self.__flops__
+        return self.__flops__
+    else:
+        sum = 0
+        for m in self.children():
+            sum += m.accumulate_flops(layer_flops)
+        layer_flops[self] = sum
+        return sum
+
+
+def get_model_parameters_number(model):
+    params_num = sum(p.numel() for p in model.parameters())
+    return params_num
+
+
+def add_flops_counting_methods(net_main_module):
+    # adding additional methods to the existing module object,
+    # this is done this way so that each function has access to self object
+    net_main_module.start_flops_count = start_flops_count.__get__(net_main_module)
+    net_main_module.stop_flops_count = stop_flops_count.__get__(net_main_module)
+    net_main_module.reset_flops_count = reset_flops_count.__get__(net_main_module)
+    net_main_module.compute_average_flops_cost = compute_average_flops_cost.__get__(
+                                                    net_main_module)
+
+    net_main_module.reset_flops_count()
+
+    return net_main_module
+
+def compute_average_flops_cost(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+    Returns current mean flops consumption per image.
+    """
+
+    for m in self.modules():
+        m.accumulate_flops = accumulate_flops.__get__(m)
+
+    layer_flops = {}
+    flops_sum = self.accumulate_flops(layer_flops)
+
+    for m in self.modules():
+        if hasattr(m, 'accumulate_flops'):
+            del m.accumulate_flops
+
+    layer_params = {}
+    for m in self.modules():
+        layer_params[m] = get_model_parameters_number(m)
+
+    params_sum = get_model_parameters_number(self)
+    return flops_sum / self.__batch_counter__, params_sum, layer_flops, layer_params
+
+
+def start_flops_count(self, **kwargs):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+    Activates the computation of mean flops consumption per image.
+    Call it before you run the network.
+    """
+    add_batch_counter_hook_function(self)
+
+    seen_types = set()
+
+    def add_flops_counter_hook_function(module, ost, verbose, ignore_list):
+        if type(module) in ignore_list:
+            seen_types.add(type(module))
+            if is_supported_instance(module):
+                module.__params__ = 0
+        elif is_supported_instance(module):
+            if hasattr(module, '__flops_handle__'):
+                return
+            if type(module) in CUSTOM_MODULES_MAPPING:
+                handle = module.register_forward_hook(
+                                        CUSTOM_MODULES_MAPPING[type(module)])
+            else:
+                handle = module.register_forward_hook(MODULES_MAPPING[type(module)])
+            module.__flops_handle__ = handle
+            seen_types.add(type(module))
+        else:
+            if verbose and not type(module) in (nn.Sequential, nn.ModuleList) and \
+               not type(module) in seen_types:
+                print('Warning: module ' + type(module).__name__ +
+                      ' is treated as a zero-op.', file=ost)
+            seen_types.add(type(module))
+
+    self.apply(partial(add_flops_counter_hook_function, **kwargs))
+
+
+def stop_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+    Stops computing the mean flops consumption per image.
+    Call whenever you want to pause the computation.
+    """
+    remove_batch_counter_hook_function(self)
+    self.apply(remove_flops_counter_hook_function)
+    self.apply(remove_flops_counter_variables)
+
+
+def reset_flops_count(self):
+    """
+    A method that will be available after add_flops_counting_methods() is called
+    on a desired net object.
+    Resets statistics computed so far.
+    """
+    add_batch_counter_variables_or_reset(self)
+    self.apply(add_flops_counter_variable_or_reset)
+
+
+# ---- Internal functions
+def batch_counter_hook(module, input, output):
+    batch_size = 1
+    if len(input) > 0:
+        # Can have multiple inputs, getting the first one
+        input = input[0]
+        batch_size = len(input)
+    else:
+        pass
+        print('Warning! No positional inputs found for a module,'
+              ' assuming batch size is 1.')
+    module.__batch_counter__ += batch_size
+
+
+def add_batch_counter_variables_or_reset(module):
+
+    module.__batch_counter__ = 0
+
+
+def add_batch_counter_hook_function(module):
+    if hasattr(module, '__batch_counter_handle__'):
+        return
+
+    handle = module.register_forward_hook(batch_counter_hook)
+    module.__batch_counter_handle__ = handle
+
+
+def remove_batch_counter_hook_function(module):
+    if hasattr(module, '__batch_counter_handle__'):
+        module.__batch_counter_handle__.remove()
+        del module.__batch_counter_handle__
+
+
+def add_flops_counter_variable_or_reset(module):
+    if is_supported_instance(module):
+        if hasattr(module, '__flops__') or hasattr(module, '__params__'):
+            print('Warning: variables __flops__ or __params__ are already '
+                  'defined for the module' + type(module).__name__ +
+                  ' ptflops can affect your code!')
+            module.__ptflops_backup_flops__ = module.__flops__
+            module.__ptflops_backup_params__ = module.__params__
+        module.__flops__ = 0
+        module.__params__ = get_model_parameters_number(module)
+
+
+def is_supported_instance(module):
+    if type(module) in MODULES_MAPPING or type(module) in CUSTOM_MODULES_MAPPING:
+        return True
+    return False
+
+
+def remove_flops_counter_hook_function(module):
+    if is_supported_instance(module):
+        if hasattr(module, '__flops_handle__'):
+            module.__flops_handle__.remove()
+            del module.__flops_handle__
+
+
+def remove_flops_counter_variables(module):
+    if is_supported_instance(module):
+        if hasattr(module, '__flops__'):
+            del module.__flops__
+            if hasattr(module, '__ptflops_backup_flops__'):
+                module.__flops__ = module.__ptflops_backup_flops__
+        if hasattr(module, '__params__'):
+            del module.__params__
+            if hasattr(module, '__ptflops_backup_params__'):
+                module.__params__ = module.__ptflops_backup_params__
+    

ixformer_sdk/contrib/DeepCache/sd/pipeline_stable_diffusion.py

@@ -0,0 +1,812 @@
+# Copyright 2023 The HuggingFace 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.
+import time 
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+import numpy as np
+from packaging import version
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+
+from diffusers.configuration_utils import FrozenDict
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL
+from diffusers.models.lora import adjust_lora_scale_text_encoder
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import (
+    deprecate,
+    logging,
+    replace_example_docstring,
+)
+from diffusers.utils.torch_utils import randn_tensor
+
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+
+from .unet_2d_condition import UNet2DConditionModel
+from .pipeline_utils import DiffusionPipeline
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import StableDiffusionPipeline
+
+        >>> pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
+        >>> pipe = pipe.to("cuda")
+
+        >>> prompt = "a photo of an astronaut riding a horse on mars"
+        >>> image = pipe(prompt).images[0]
+        ```
+"""
+
+def sample_gaussian_centered(n=1000, sample_size=100, std_dev=100):
+    samples = []
+    
+    while len(samples) < sample_size:
+        # Sample from a Gaussian centered at n/2
+        sample = int(np.random.normal(loc=n/2, scale=std_dev))
+        
+        # Check if the sample is in bounds
+        if 1 <= sample < n and sample not in samples:
+            samples.append(sample)
+    
+    return samples
+
+def sample_from_quad(total_numbers, n_samples, pow=1.2):
+    while pow > 1:
+        # Generate linearly spaced values between 0 and a max value
+        x_values = np.linspace(0, total_numbers**(1/pow), n_samples+1)
+
+        # Raise these values to the power of 1.5 to get a non-linear distribution
+        indices = np.unique(np.int32(x_values**pow))[:-1]
+        if len(indices) == n_samples:
+            break
+        pow -=0.02
+    if pow <= 1:
+        raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.")
+    return indices, pow
+
+def sample_from_quad_center(total_numbers, n_samples, center, pow=1.2):
+    while pow > 1:
+        # Generate linearly spaced values between 0 and a max value
+        x_values = np.linspace((-center)**(1/pow), (total_numbers-center)**(1/pow), n_samples+1)
+        indices = [0] + [x+center for x in np.unique(np.int32(x_values**pow))[1:-1]]
+        if len(indices) == n_samples:
+            break
+        pow -=0.02
+    if pow <= 1:
+        raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.")
+    return indices, pow
+
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    """
+    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
+    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
+    """
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    # rescale the results from guidance (fixes overexposure)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
+
+class StableDiffusionPipeline(DiffusionPipeline, TextualInversionLoaderMixin, LoraLoaderMixin, FromSingleFileMixin):
+    r"""
+    Pipeline for text-to-image generation using Stable Diffusion.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+
+    The pipeline also inherits the following loading methods:
+        - [`~loaders.TextualInversionLoaderMixin.load_textual_inversion`] for loading textual inversion embeddings
+        - [`~loaders.LoraLoaderMixin.load_lora_weights`] for loading LoRA weights
+        - [`~loaders.LoraLoaderMixin.save_lora_weights`] for saving LoRA weights
+        - [`~loaders.FromSingleFileMixin.from_single_file`] for loading `.ckpt` files
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
+        text_encoder ([`~transformers.CLIPTextModel`]):
+            Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
+        tokenizer ([`~transformers.CLIPTokenizer`]):
+            A `CLIPTokenizer` to tokenize text.
+        unet ([`UNet2DConditionModel`]):
+            A `UNet2DConditionModel` to denoise the encoded image latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
+            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
+        safety_checker ([`StableDiffusionSafetyChecker`]):
+            Classification module that estimates whether generated images could be considered offensive or harmful.
+            Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
+            about a model's potential harms.
+        feature_extractor ([`~transformers.CLIPImageProcessor`]):
+            A `CLIPImageProcessor` to extract features from generated images; used as inputs to the `safety_checker`.
+    """
+    model_cpu_offload_seq = "text_encoder->unet->vae"
+    _optional_components = ["safety_checker", "feature_extractor"]
+    _exclude_from_cpu_offload = ["safety_checker"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        safety_checker: StableDiffusionSafetyChecker,
+        feature_extractor: CLIPImageProcessor,
+        requires_safety_checker: bool = True,
+    ):
+        super().__init__()
+
+        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
+                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
+                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
+                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
+                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
+                " file"
+            )
+            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["steps_offset"] = 1
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
+                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
+                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
+                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
+                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
+            )
+            deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["clip_sample"] = False
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if safety_checker is None and requires_safety_checker:
+            logger.warning(
+                f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
+                " that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
+                " results in services or applications open to the public. Both the diffusers team and Hugging Face"
+                " strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
+                " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
+                " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
+            )
+
+        if safety_checker is not None and feature_extractor is None:
+            raise ValueError(
+                "Make sure to define a feature extractor when loading {self.__class__} if you want to use the safety"
+                " checker. If you do not want to use the safety checker, you can pass `'safety_checker=None'` instead."
+            )
+
+        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(
+            version.parse(unet.config._diffusers_version).base_version
+        ) < version.parse("0.9.0.dev0")
+        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
+        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
+            deprecation_message = (
+                "The configuration file of the unet has set the default `sample_size` to smaller than"
+                " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"
+                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
+                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
+                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
+                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
+                " in the config might lead to incorrect results in future versions. If you have downloaded this"
+                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
+                " the `unet/config.json` file"
+            )
+            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(unet.config)
+            new_config["sample_size"] = 64
+            unet._internal_dict = FrozenDict(new_config)
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.register_to_config(requires_safety_checker=requires_safety_checker)
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
+        compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+        compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+        processing larger images.
+        """
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        lora_scale: Optional[float] = None,
+    ):
+        deprecation_message = "`_encode_prompt()` is deprecated and it will be removed in a future version. Use `encode_prompt()` instead. Also, be aware that the output format changed from a concatenated tensor to a tuple."
+        deprecate("_encode_prompt()", "1.0.0", deprecation_message, standard_warn=False)
+
+        prompt_embeds_tuple = self.encode_prompt(
+            prompt=prompt,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=lora_scale,
+        )
+
+        # concatenate for backwards comp
+        prompt_embeds = torch.cat([prompt_embeds_tuple[1], prompt_embeds_tuple[0]])
+
+        return prompt_embeds
+
+    def encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        lora_scale: Optional[float] = None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            lora_scale (`float`, *optional*):
+                A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+        """
+        # set lora scale so that monkey patched LoRA
+        # function of text encoder can correctly access it
+        if lora_scale is not None and isinstance(self, LoraLoaderMixin):
+            self._lora_scale = lora_scale
+
+            # dynamically adjust the LoRA scale
+            adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=self.tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = self.tokenizer.batch_decode(
+                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
+                )
+                logger.warning(
+                    "The following part of your input was truncated because CLIP can only handle sequences up to"
+                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"
+                )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = text_inputs.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            prompt_embeds = self.text_encoder(
+                text_input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            prompt_embeds = prompt_embeds[0]
+
+        if self.text_encoder is not None:
+            prompt_embeds_dtype = self.text_encoder.dtype
+        elif self.unet is not None:
+            prompt_embeds_dtype = self.unet.dtype
+        else:
+            prompt_embeds_dtype = prompt_embeds.dtype
+
+        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif isinstance(negative_prompt, str):
+                uncond_tokens = [negative_prompt]
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+            else:
+                uncond_tokens = negative_prompt
+
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
+
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+        return prompt_embeds, negative_prompt_embeds
+
+    def run_safety_checker(self, image, device, dtype):
+        if self.safety_checker is None:
+            has_nsfw_concept = None
+        else:
+            if torch.is_tensor(image):
+                feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
+            else:
+                feature_extractor_input = self.image_processor.numpy_to_pil(image)
+            safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
+            image, has_nsfw_concept = self.safety_checker(
+                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
+            )
+        return image, has_nsfw_concept
+
+    def decode_latents(self, latents):
+        deprecation_message = "The decode_latents method is deprecated and will be removed in 1.0.0. Please use VaeImageProcessor.postprocess(...) instead"
+        deprecate("decode_latents", "1.0.0", deprecation_message, standard_warn=False)
+
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents, return_dict=False)[0]
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        callback_steps,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+    ):
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        cache_interval: int = 1,
+        cache_layer_id: int = None,
+        cache_block_id: int = None,
+        uniform: bool = True,
+        pow: float = None,
+        center: int = None,
+        output_all_sequence: bool = False,
+    ):
+        r"""
+        The call function to the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
+            height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                A higher guidance scale value encourages the model to generate images closely linked to the text
+                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide what to not include in image generation. If not defined, you need to
+                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
+                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor is generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
+                provided, text embeddings are generated from the `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
+                not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
+                plain tuple.
+            callback (`Callable`, *optional*):
+                A function that calls every `callback_steps` steps during inference. The function is called with the
+                following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function is called. If not specified, the callback is called at
+                every step.
+            cross_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
+                [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            guidance_rescale (`float`, *optional*, defaults to 0.7):
+                Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are
+                Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when
+                using zero terminal SNR.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
+                otherwise a `tuple` is returned where the first element is a list with the generated images and the
+                second element is a list of `bool`s indicating whether the corresponding generated image contains
+                "not-safe-for-work" (nsfw) content.
+        """
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds
+        )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
+        )
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=text_encoder_lora_scale,
+        )
+        # For classifier free guidance, we need to do two forward passes.
+        # Here we concatenate the unconditional and text embeddings into a single batch
+        # to avoid doing two forward passes
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        prv_features = None
+        latents_list = [latents]
+
+        if cache_interval == 1:
+            interval_seq = list(range(num_inference_steps))
+        else:
+            if uniform:
+                interval_seq = list(range(0, num_inference_steps, cache_interval))
+            else:
+                num_slow_step = num_inference_steps//cache_interval
+                if num_inference_steps%cache_interval != 0:
+                    num_slow_step += 1
+                
+                interval_seq, pow = sample_from_quad_center(num_inference_steps, num_slow_step, center=center, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,]
+                #interval_seq, pow = sample_from_quad(num_inference_steps, num_inference_steps//cache_interval, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,]
+        
+        interval_seq = sorted(interval_seq)
+        #print(interval_seq, len(interval_seq), pow)
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            #print("[INFO] Update Feature Interval = {}, Update Layer Number = {}, Update Block Number = {}".format(cache_interval, cache_layer_id, cache_block_id))
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                if i in interval_seq:
+                    prv_features = None            
+
+                # predict the noise residual
+                noise_pred, prv_features = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    replicate_prv_feature=prv_features,
+                    quick_replicate= cache_interval>1,
+                    cache_layer_id=cache_layer_id,
+                    cache_block_id=cache_block_id,
+                    return_dict=False,
+                )
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                if do_classifier_free_guidance and guidance_rescale > 0.0:
+                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                latents_list.append(latents)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+        if not output_type == "latent":
+            if output_all_sequence:
+                image = [self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] for latents in latents_list]
+                has_nsfw_concept = None #self.run_safety_checker(images[0], device, prompt_embeds.dtype)
+                num_img = len(image)
+            else:
+                image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+                has_nsfw_concept = None
+                num_img = image.shape[0]
+        else:
+            image = latents
+            has_nsfw_concept = None
+
+        if has_nsfw_concept is None:
+            do_denormalize = [True] * num_img
+        else:
+            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
+
+        if output_all_sequence:
+            image = [self.image_processor.postprocess(img, output_type=output_type, do_denormalize=do_denormalize) for img in image]
+        else:
+            image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+        if not return_dict:
+            return (image, has_nsfw_concept,)
+
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)

ixformer_sdk/contrib/DeepCache/sd/pipeline_text_to_video_zero.py

@@ -0,0 +1,741 @@
+import copy
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Union
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn.functional as F
+from torch.nn.functional import grid_sample
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+
+from diffusers.models import AutoencoderKL
+from .unet_2d_condition import UNet2DConditionModel
+from .pipeline_stable_diffusion import StableDiffusionPipeline, StableDiffusionSafetyChecker
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+def sample_gaussian_centered(n=1000, sample_size=100, std_dev=100):
+    samples = []
+    
+    while len(samples) < sample_size:
+        # Sample from a Gaussian centered at n/2
+        sample = int(np.random.normal(loc=n/2, scale=std_dev))
+        
+        # Check if the sample is in bounds
+        if 1 <= sample < n and sample not in samples:
+            samples.append(sample)
+    
+    return samples
+
+def sample_from_quad(total_numbers, n_samples, pow=1.2):
+    while pow > 1:
+        # Generate linearly spaced values between 0 and a max value
+        x_values = np.linspace(0, total_numbers**(1/pow), n_samples+1)
+
+        # Raise these values to the power of 1.5 to get a non-linear distribution
+        indices = np.unique(np.int32(x_values**pow))[:-1]
+        if len(indices) == n_samples:
+            break
+        pow -=0.02
+    if pow <= 1:
+        raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.")
+    return indices, pow
+
+def sample_from_quad_center(total_numbers, n_samples, center, pow=1.2):
+    while pow > 1:
+        # Generate linearly spaced values between 0 and a max value
+        x_values = np.linspace((-center)**(1/pow), (total_numbers-center)**(1/pow), n_samples+1)
+        indices = [0] + [x+center for x in np.unique(np.int32(x_values**pow))[1:-1]]
+        if len(indices) == n_samples:
+            break
+        pow -=0.02
+    if pow <= 1:
+        raise ValueError("Cannot find suitable pow. Please adjust n_samples or decrease center.")
+    return indices, pow
+
+def rearrange_0(tensor, f):
+    F, C, H, W = tensor.size()
+    tensor = torch.permute(torch.reshape(tensor, (F // f, f, C, H, W)), (0, 2, 1, 3, 4))
+    return tensor
+
+
+def rearrange_1(tensor):
+    B, C, F, H, W = tensor.size()
+    return torch.reshape(torch.permute(tensor, (0, 2, 1, 3, 4)), (B * F, C, H, W))
+
+
+def rearrange_3(tensor, f):
+    F, D, C = tensor.size()
+    return torch.reshape(tensor, (F // f, f, D, C))
+
+
+def rearrange_4(tensor):
+    B, F, D, C = tensor.size()
+    return torch.reshape(tensor, (B * F, D, C))
+
+
+class CrossFrameAttnProcessor:
+    """
+    Cross frame attention processor. Each frame attends the first frame.
+
+    Args:
+        batch_size: The number that represents actual batch size, other than the frames.
+            For example, calling unet with a single prompt and num_images_per_prompt=1, batch_size should be equal to
+            2, due to classifier-free guidance.
+    """
+
+    def __init__(self, batch_size=2):
+        self.batch_size = batch_size
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None):
+        batch_size, sequence_length, _ = hidden_states.shape
+        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+        query = attn.to_q(hidden_states)
+
+        is_cross_attention = encoder_hidden_states is not None
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        # Cross Frame Attention
+        if not is_cross_attention:
+            video_length = key.size()[0] // self.batch_size
+            first_frame_index = [0] * video_length
+
+            # rearrange keys to have batch and frames in the 1st and 2nd dims respectively
+            key = rearrange_3(key, video_length)
+            key = key[:, first_frame_index]
+            # rearrange values to have batch and frames in the 1st and 2nd dims respectively
+            value = rearrange_3(value, video_length)
+            value = value[:, first_frame_index]
+
+            # rearrange back to original shape
+            key = rearrange_4(key)
+            value = rearrange_4(value)
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states
+
+
+class CrossFrameAttnProcessor2_0:
+    """
+    Cross frame attention processor with scaled_dot_product attention of Pytorch 2.0.
+
+    Args:
+        batch_size: The number that represents actual batch size, other than the frames.
+            For example, calling unet with a single prompt and num_images_per_prompt=1, batch_size should be equal to
+            2, due to classifier-free guidance.
+    """
+
+    def __init__(self, batch_size=2):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+        self.batch_size = batch_size
+
+    def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None):
+        batch_size, sequence_length, _ = (
+            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+        )
+        inner_dim = hidden_states.shape[-1]
+
+        if attention_mask is not None:
+            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+            # scaled_dot_product_attention expects attention_mask shape to be
+            # (batch, heads, source_length, target_length)
+            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+        query = attn.to_q(hidden_states)
+
+        is_cross_attention = encoder_hidden_states is not None
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        elif attn.norm_cross:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        # Cross Frame Attention
+        if not is_cross_attention:
+            video_length = max(1, key.size()[0] // self.batch_size)
+            first_frame_index = [0] * video_length
+
+            # rearrange keys to have batch and frames in the 1st and 2nd dims respectively
+            key = rearrange_3(key, video_length)
+            key = key[:, first_frame_index]
+            # rearrange values to have batch and frames in the 1st and 2nd dims respectively
+            value = rearrange_3(value, video_length)
+            value = value[:, first_frame_index]
+
+            # rearrange back to original shape
+            key = rearrange_4(key)
+            value = rearrange_4(value)
+
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # the output of sdp = (batch, num_heads, seq_len, head_dim)
+        # TODO: add support for attn.scale when we move to Torch 2.1
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+        return hidden_states
+
+
+@dataclass
+class TextToVideoPipelineOutput(BaseOutput):
+    r"""
+    Output class for zero-shot text-to-video pipeline.
+
+    Args:
+        images (`[List[PIL.Image.Image]`, `np.ndarray`]):
+            List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+            num_channels)`.
+        nsfw_content_detected (`[List[bool]]`):
+            List indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content or
+            `None` if safety checking could not be performed.
+    """
+
+    images: Union[List[PIL.Image.Image], np.ndarray]
+    nsfw_content_detected: Optional[List[bool]]
+
+
+def coords_grid(batch, ht, wd, device):
+    # Adapted from https://github.com/princeton-vl/RAFT/blob/master/core/utils/utils.py
+    coords = torch.meshgrid(torch.arange(ht, device=device), torch.arange(wd, device=device))
+    coords = torch.stack(coords[::-1], dim=0).float()
+    return coords[None].repeat(batch, 1, 1, 1)
+
+
+def warp_single_latent(latent, reference_flow):
+    """
+    Warp latent of a single frame with given flow
+
+    Args:
+        latent: latent code of a single frame
+        reference_flow: flow which to warp the latent with
+
+    Returns:
+        warped: warped latent
+    """
+    _, _, H, W = reference_flow.size()
+    _, _, h, w = latent.size()
+    coords0 = coords_grid(1, H, W, device=latent.device).to(latent.dtype)
+
+    coords_t0 = coords0 + reference_flow
+    coords_t0[:, 0] /= W
+    coords_t0[:, 1] /= H
+
+    coords_t0 = coords_t0 * 2.0 - 1.0
+    coords_t0 = F.interpolate(coords_t0, size=(h, w), mode="bilinear")
+    coords_t0 = torch.permute(coords_t0, (0, 2, 3, 1))
+
+    warped = grid_sample(latent, coords_t0, mode="nearest", padding_mode="reflection")
+    return warped
+
+
+def create_motion_field(motion_field_strength_x, motion_field_strength_y, frame_ids, device, dtype):
+    """
+    Create translation motion field
+
+    Args:
+        motion_field_strength_x: motion strength along x-axis
+        motion_field_strength_y: motion strength along y-axis
+        frame_ids: indexes of the frames the latents of which are being processed.
+            This is needed when we perform chunk-by-chunk inference
+        device: device
+        dtype: dtype
+
+    Returns:
+
+    """
+    seq_length = len(frame_ids)
+    reference_flow = torch.zeros((seq_length, 2, 512, 512), device=device, dtype=dtype)
+    for fr_idx in range(seq_length):
+        reference_flow[fr_idx, 0, :, :] = motion_field_strength_x * (frame_ids[fr_idx])
+        reference_flow[fr_idx, 1, :, :] = motion_field_strength_y * (frame_ids[fr_idx])
+    return reference_flow
+
+
+def create_motion_field_and_warp_latents(motion_field_strength_x, motion_field_strength_y, frame_ids, latents):
+    """
+    Creates translation motion and warps the latents accordingly
+
+    Args:
+        motion_field_strength_x: motion strength along x-axis
+        motion_field_strength_y: motion strength along y-axis
+        frame_ids: indexes of the frames the latents of which are being processed.
+            This is needed when we perform chunk-by-chunk inference
+        latents: latent codes of frames
+
+    Returns:
+        warped_latents: warped latents
+    """
+    motion_field = create_motion_field(
+        motion_field_strength_x=motion_field_strength_x,
+        motion_field_strength_y=motion_field_strength_y,
+        frame_ids=frame_ids,
+        device=latents.device,
+        dtype=latents.dtype,
+    )
+    warped_latents = latents.clone().detach()
+    for i in range(len(warped_latents)):
+        warped_latents[i] = warp_single_latent(latents[i][None], motion_field[i][None])
+    return warped_latents
+
+
+class TextToVideoZeroPipeline(StableDiffusionPipeline):
+    r"""
+    Pipeline for zero-shot text-to-video generation using Stable Diffusion.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            Frozen text-encoder ([clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)).
+        tokenizer (`CLIPTokenizer`):
+            A [`~transformers.CLIPTokenizer`] to tokenize text.
+        unet ([`UNet2DConditionModel`]):
+            A [`UNet3DConditionModel`] to denoise the encoded video latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
+            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
+        safety_checker ([`StableDiffusionSafetyChecker`]):
+            Classification module that estimates whether generated images could be considered offensive or harmful.
+            Please refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for more details
+            about a model's potential harms.
+        feature_extractor ([`CLIPImageProcessor`]):
+            A [`CLIPImageProcessor`] to extract features from generated images; used as inputs to the `safety_checker`.
+    """
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        safety_checker: StableDiffusionSafetyChecker,
+        feature_extractor: CLIPImageProcessor,
+        requires_safety_checker: bool = True,
+    ):
+        super().__init__(
+            vae, text_encoder, tokenizer, unet, scheduler, safety_checker, feature_extractor, requires_safety_checker
+        )
+        processor = (
+            CrossFrameAttnProcessor2_0(batch_size=2)
+            if hasattr(F, "scaled_dot_product_attention")
+            else CrossFrameAttnProcessor(batch_size=2)
+        )
+        self.unet.set_attn_processor(processor)
+
+    def forward_loop(self, x_t0, t0, t1, generator):
+        """
+        Perform DDPM forward process from time t0 to t1. This is the same as adding noise with corresponding variance.
+
+        Args:
+            x_t0:
+                Latent code at time t0.
+            t0:
+                Timestep at t0.
+            t1:
+                Timestamp at t1.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+
+        Returns:
+            x_t1:
+                Forward process applied to x_t0 from time t0 to t1.
+        """
+        eps = randn_tensor(x_t0.size(), generator=generator, dtype=x_t0.dtype, device=x_t0.device)
+        alpha_vec = torch.prod(self.scheduler.alphas[t0:t1])
+        x_t1 = torch.sqrt(alpha_vec) * x_t0 + torch.sqrt(1 - alpha_vec) * eps
+        return x_t1
+
+    def backward_loop(
+        self,
+        latents,
+        timesteps,
+        prompt_embeds,
+        guidance_scale,
+        callback,
+        callback_steps,
+        num_warmup_steps,
+        extra_step_kwargs,
+        prv_features,
+        interval_seq,
+        cache_interval,
+        cache_block_id,
+        cache_layer_id,
+        cross_attention_kwargs=None,
+    ):
+        """
+        Perform backward process given list of time steps.
+
+        Args:
+            latents:
+                Latents at time timesteps[0].
+            timesteps:
+                Time steps along which to perform backward process.
+            prompt_embeds:
+                Pre-generated text embeddings.
+            guidance_scale:
+                A higher guidance scale value encourages the model to generate images closely linked to the text
+                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
+            callback (`Callable`, *optional*):
+                A function that calls every `callback_steps` steps during inference. The function is called with the
+                following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function is called. If not specified, the callback is called at
+                every step.
+            extra_step_kwargs:
+                Extra_step_kwargs.
+            cross_attention_kwargs:
+                A kwargs dictionary that if specified is passed along to the [`AttentionProcessor`] as defined in
+                [`self.processor`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            num_warmup_steps:
+                number of warmup steps.
+
+        Returns:
+            latents:
+                Latents of backward process output at time timesteps[-1].
+        """
+        do_classifier_free_guidance = guidance_scale > 1.0
+        num_steps = (len(timesteps) - num_warmup_steps) // self.scheduler.order
+        with self.progress_bar(total=num_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+                ########
+                if i in interval_seq:
+                    prv_features = None
+                # predict the noise residual
+                noise_pred, prv_features = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    replicate_prv_feature=prv_features,
+                    quick_replicate= cache_interval>1,
+                    cache_layer_id=cache_layer_id,
+                    cache_block_id=cache_block_id,
+                    return_dict=False,
+                )
+                ########
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        callback(step_idx, t, latents)
+        return latents.clone().detach()
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        video_length: Optional[int] = 8,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_videos_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        motion_field_strength_x: float = 12,
+        motion_field_strength_y: float = 12,
+        output_type: Optional[str] = "tensor",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: Optional[int] = 1,
+        t0: int = 44,
+        t1: int = 47,
+        frame_ids: Optional[List[int]] = None,
+        ########
+        cache_interval: int = 1,
+        cache_layer_id: int = None,
+        cache_block_id: int = None,
+        uniform: bool = True,
+        pow: float = None,
+        center: int = None,
+        output_all_sequence: bool = False,
+        ########
+    ):
+        """
+        The call function to the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
+            video_length (`int`, *optional*, defaults to 8):
+                The number of generated video frames.
+            height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                A higher guidance scale value encourages the model to generate images closely linked to the text
+                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide what to not include in video generation. If not defined, you need to
+                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of videos to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
+                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for video
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor is generated by sampling using the supplied random `generator`.
+            output_type (`str`, *optional*, defaults to `"numpy"`):
+                The output format of the generated video. Choose between `"latent"` and `"numpy"`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a
+                [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoPipelineOutput`] instead of
+                a plain tuple.
+            callback (`Callable`, *optional*):
+                A function that calls every `callback_steps` steps during inference. The function is called with the
+                following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function is called. If not specified, the callback is called at
+                every step.
+            motion_field_strength_x (`float`, *optional*, defaults to 12):
+                Strength of motion in generated video along x-axis. See the [paper](https://arxiv.org/abs/2303.13439),
+                Sect. 3.3.1.
+            motion_field_strength_y (`float`, *optional*, defaults to 12):
+                Strength of motion in generated video along y-axis. See the [paper](https://arxiv.org/abs/2303.13439),
+                Sect. 3.3.1.
+            t0 (`int`, *optional*, defaults to 44):
+                Timestep t0. Should be in the range [0, num_inference_steps - 1]. See the
+                [paper](https://arxiv.org/abs/2303.13439), Sect. 3.3.1.
+            t1 (`int`, *optional*, defaults to 47):
+                Timestep t0. Should be in the range [t0 + 1, num_inference_steps - 1]. See the
+                [paper](https://arxiv.org/abs/2303.13439), Sect. 3.3.1.
+            frame_ids (`List[int]`, *optional*):
+                Indexes of the frames that are being generated. This is used when generating longer videos
+                chunk-by-chunk.
+
+        Returns:
+            [`~pipelines.text_to_video_synthesis.pipeline_text_to_video_zero.TextToVideoPipelineOutput`]:
+                The output contains a `ndarray` of the generated video, when `output_type` != `"latent"`, otherwise a
+                latent code of generated videos and a list of `bool`s indicating whether the corresponding generated
+                video contains "not-safe-for-work" (nsfw) content..
+        """
+        assert video_length > 0
+        if frame_ids is None:
+            frame_ids = list(range(video_length))
+        assert len(frame_ids) == video_length
+
+        assert num_videos_per_prompt == 1
+
+        if isinstance(prompt, str):
+            prompt = [prompt]
+        if isinstance(negative_prompt, str):
+            negative_prompt = [negative_prompt]
+
+        # Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # Check inputs. Raise error if not correct
+        self.check_inputs(prompt, height, width, callback_steps)
+
+        # Define call parameters
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # Encode input prompt
+        prompt_embeds = self._encode_prompt(
+            prompt, device, num_videos_per_prompt, do_classifier_free_guidance, negative_prompt
+        )
+
+        # Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_videos_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        # Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        prv_features = None #record cache feature ****
+        latents_list = [latents] 
+
+        if cache_interval == 1:
+            interval_seq = list(range(num_inference_steps))
+        else:
+            if uniform:
+                interval_seq = list(range(0, num_inference_steps, cache_interval))
+            else:
+                num_slow_step = num_inference_steps//cache_interval
+                if num_inference_steps%cache_interval != 0:
+                    num_slow_step += 1
+                
+                interval_seq, pow = sample_from_quad_center(num_inference_steps, num_slow_step, center=center, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,]
+                #interval_seq, pow = sample_from_quad(num_inference_steps, num_inference_steps//cache_interval, pow=pow)#[0, 3, 6, 9, 12, 16, 22, 28, 35, 43,]
+        
+        interval_seq = sorted(interval_seq)
+
+        # Perform the first backward process up to time T_1
+        x_1_t1 = self.backward_loop(
+            timesteps=timesteps[: -t1 - 1],
+            prompt_embeds=prompt_embeds,
+            latents=latents,
+            guidance_scale=guidance_scale,
+            callback=callback,
+            callback_steps=callback_steps,
+            extra_step_kwargs=extra_step_kwargs,
+            num_warmup_steps=num_warmup_steps,
+            prv_features=prv_features,
+            interval_seq=interval_seq,
+            cache_interval=cache_interval,
+            cache_block_id=cache_block_id,
+            cache_layer_id=cache_layer_id,
+        )
+        scheduler_copy = copy.deepcopy(self.scheduler)
+
+        # Perform the second backward process up to time T_0
+        x_1_t0 = self.backward_loop(
+            timesteps=timesteps[-t1 - 1 : -t0 - 1],
+            prompt_embeds=prompt_embeds,
+            latents=x_1_t1,
+            guidance_scale=guidance_scale,
+            callback=callback,
+            callback_steps=callback_steps,
+            extra_step_kwargs=extra_step_kwargs,
+            num_warmup_steps=0,
+            prv_features=prv_features,
+            interval_seq=interval_seq,
+            cache_interval=cache_interval,
+            cache_block_id=cache_block_id,
+            cache_layer_id=cache_layer_id,
+        )
+
+        # Propagate first frame latents at time T_0 to remaining frames
+        x_2k_t0 = x_1_t0.repeat(video_length - 1, 1, 1, 1)
+
+        # Add motion in latents at time T_0
+        x_2k_t0 = create_motion_field_and_warp_latents(
+            motion_field_strength_x=motion_field_strength_x,
+            motion_field_strength_y=motion_field_strength_y,
+            latents=x_2k_t0,
+            frame_ids=frame_ids[1:],
+        )
+
+        # Perform forward process up to time T_1
+        x_2k_t1 = self.forward_loop(
+            x_t0=x_2k_t0,
+            t0=timesteps[-t0 - 1].item(),
+            t1=timesteps[-t1 - 1].item(),
+            generator=generator,
+        )
+
+        # Perform backward process from time T_1 to 0
+        x_1k_t1 = torch.cat([x_1_t1, x_2k_t1])
+        b, l, d = prompt_embeds.size()
+        prompt_embeds = prompt_embeds[:, None].repeat(1, video_length, 1, 1).reshape(b * video_length, l, d)
+
+        self.scheduler = scheduler_copy
+        x_1k_0 = self.backward_loop(
+            timesteps=timesteps[-t1 - 1 :],
+            prompt_embeds=prompt_embeds,
+            latents=x_1k_t1,
+            guidance_scale=guidance_scale,
+            callback=callback,
+            callback_steps=callback_steps,
+            extra_step_kwargs=extra_step_kwargs,
+            num_warmup_steps=0,
+            prv_features=prv_features,
+            interval_seq=interval_seq,
+            cache_interval=cache_interval,
+            cache_block_id=cache_block_id,
+            cache_layer_id=cache_layer_id,
+        )
+        latents = x_1k_0
+
+        # manually for max memory savings
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.unet.to("cpu")
+        torch.cuda.empty_cache()
+
+        if output_type == "latent":
+            image = latents
+            has_nsfw_concept = None
+        else:
+            image = self.decode_latents(latents)
+            # Run safety checker
+            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return TextToVideoPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)

ixformer_sdk/contrib/DeepCache/svd/pipeline_stable_video_diffusion.py

@@ -0,0 +1,659 @@
+# Copyright 2023 The HuggingFace 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.
+
+import inspect
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Optional, Union
+
+import numpy as np
+import PIL.Image
+import torch
+from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models import AutoencoderKLTemporalDecoder, UNetSpatioTemporalConditionModel
+from diffusers.schedulers import EulerDiscreteScheduler
+from diffusers.utils import BaseOutput, logging
+from diffusers.utils.torch_utils import randn_tensor
+from .pipeline_utils import DiffusionPipeline
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def _append_dims(x, target_dims):
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}, which is less")
+    return x[(...,) + (None,) * dims_to_append]
+
+
+def tensor2vid(video: torch.Tensor, processor, output_type="np"):
+    # Based on:
+    # https://github.com/modelscope/modelscope/blob/1509fdb973e5871f37148a4b5e5964cafd43e64d/modelscope/pipelines/multi_modal/text_to_video_synthesis_pipeline.py#L78
+
+    batch_size, channels, num_frames, height, width = video.shape
+    outputs = []
+    for batch_idx in range(batch_size):
+        batch_vid = video[batch_idx].permute(1, 0, 2, 3)
+        batch_output = processor.postprocess(batch_vid, output_type)
+
+        outputs.append(batch_output)
+
+    return outputs
+
+
+@dataclass
+class StableVideoDiffusionPipelineOutput(BaseOutput):
+    r"""
+    Output class for zero-shot text-to-video pipeline.
+
+    Args:
+        frames (`[List[PIL.Image.Image]`, `np.ndarray`]):
+            List of denoised PIL images of length `batch_size` or NumPy array of shape `(batch_size, height, width,
+            num_channels)`.
+    """
+
+    frames: Union[List[PIL.Image.Image], np.ndarray]
+
+
+class StableVideoDiffusionPipeline(DiffusionPipeline):
+    r"""
+    Pipeline to generate video from an input image using Stable Video Diffusion.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) model to encode and decode images to and from latent representations.
+        image_encoder ([`~transformers.CLIPVisionModelWithProjection`]):
+            Frozen CLIP image-encoder ([laion/CLIP-ViT-H-14-laion2B-s32B-b79K](https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K)).
+        unet ([`UNetSpatioTemporalConditionModel`]):cache_interval=5, cache_branch=0,
+            A `UNetSpatioTemporalConditionModel` to denoise the encoded image latents.
+        scheduler ([`EulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents.
+        feature_extractor ([`~transformers.CLIPImageProcessor`]):
+            A `CLIPImageProcessor` to extract features from generated images.
+    """
+
+    model_cpu_offload_seq = "image_encoder->unet->vae"
+    _callback_tensor_inputs = ["latents"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKLTemporalDecoder,
+        image_encoder: CLIPVisionModelWithProjection,
+        unet: UNetSpatioTemporalConditionModel,
+        scheduler: EulerDiscreteScheduler,
+        feature_extractor: CLIPImageProcessor,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            image_encoder=image_encoder,
+            unet=unet,
+            scheduler=scheduler,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+    def _encode_image(self, image, device, num_videos_per_prompt, do_classifier_free_guidance):
+        dtype = next(self.image_encoder.parameters()).dtype
+
+        if not isinstance(image, torch.Tensor):
+            image = self.image_processor.pil_to_numpy(image)
+            image = self.image_processor.numpy_to_pt(image)
+
+            # We normalize the image before resizing to match with the original implementation.
+            # Then we unnormalize it after resizing.
+            image = image * 2.0 - 1.0
+            image = _resize_with_antialiasing(image, (224, 224))
+            image = (image + 1.0) / 2.0
+
+            # Normalize the image with for CLIP input
+            image = self.feature_extractor(
+                images=image,
+                do_normalize=True,
+                do_center_crop=False,
+                do_resize=False,
+                do_rescale=False,
+                return_tensors="pt",
+            ).pixel_values
+
+        image = image.to(device=device, dtype=dtype)
+        image_embeddings = self.image_encoder(image).image_embeds
+        image_embeddings = image_embeddings.unsqueeze(1)
+
+        # duplicate image embeddings for each generation per prompt, using mps friendly method
+        bs_embed, seq_len, _ = image_embeddings.shape
+        image_embeddings = image_embeddings.repeat(1, num_videos_per_prompt, 1)
+        image_embeddings = image_embeddings.view(bs_embed * num_videos_per_prompt, seq_len, -1)
+
+        if do_classifier_free_guidance:
+            negative_image_embeddings = torch.zeros_like(image_embeddings)
+
+            # For classifier free guidance, we need to do two forward passes.
+            # Here we concatenate the unconditional and text embeddings into a single batch
+            # to avoid doing two forward passes
+            image_embeddings = torch.cat([negative_image_embeddings, image_embeddings])
+
+        return image_embeddings
+
+    def _encode_vae_image(
+        self,
+        image: torch.Tensor,
+        device,
+        num_videos_per_prompt,
+        do_classifier_free_guidance,
+    ):
+        image = image.to(device=device)
+        image_latents = self.vae.encode(image).latent_dist.mode()
+
+        if do_classifier_free_guidance:
+            negative_image_latents = torch.zeros_like(image_latents)
+
+            # For classifier free guidance, we need to do two forward passes.
+            # Here we concatenate the unconditional and text embeddings into a single batch
+            # to avoid doing two forward passes
+            image_latents = torch.cat([negative_image_latents, image_latents])
+
+        # duplicate image_latents for each generation per prompt, using mps friendly method
+        image_latents = image_latents.repeat(num_videos_per_prompt, 1, 1, 1)
+
+        return image_latents
+
+    def _get_add_time_ids(
+        self,
+        fps,
+        motion_bucket_id,
+        noise_aug_strength,
+        dtype,
+        batch_size,
+        num_videos_per_prompt,
+        do_classifier_free_guidance,
+    ):
+        add_time_ids = [fps, motion_bucket_id, noise_aug_strength]
+
+        passed_add_embed_dim = self.unet.config.addition_time_embed_dim * len(add_time_ids)
+        expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
+
+        if expected_add_embed_dim != passed_add_embed_dim:
+            raise ValueError(
+                f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
+            )
+
+        add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
+        add_time_ids = add_time_ids.repeat(batch_size * num_videos_per_prompt, 1)
+
+        if do_classifier_free_guidance:
+            add_time_ids = torch.cat([add_time_ids, add_time_ids])
+
+        return add_time_ids
+
+    def decode_latents(self, latents, num_frames, decode_chunk_size=14):
+        # [batch, frames, channels, height, width] -> [batch*frames, channels, height, width]
+        latents = latents.flatten(0, 1)
+
+        latents = 1 / self.vae.config.scaling_factor * latents
+
+        accepts_num_frames = "num_frames" in set(inspect.signature(self.vae.forward).parameters.keys())
+
+        # decode decode_chunk_size frames at a time to avoid OOM
+        frames = []
+        for i in range(0, latents.shape[0], decode_chunk_size):
+            num_frames_in = latents[i : i + decode_chunk_size].shape[0]
+            decode_kwargs = {}
+            if accepts_num_frames:
+                # we only pass num_frames_in if it's expected
+                decode_kwargs["num_frames"] = num_frames_in
+
+            frame = self.vae.decode(latents[i : i + decode_chunk_size], **decode_kwargs).sample
+            frames.append(frame)
+        frames = torch.cat(frames, dim=0)
+
+        # [batch*frames, channels, height, width] -> [batch, channels, frames, height, width]
+        frames = frames.reshape(-1, num_frames, *frames.shape[1:]).permute(0, 2, 1, 3, 4)
+
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        frames = frames.float()
+        return frames
+
+    def check_inputs(self, image, height, width):
+        if (
+            not isinstance(image, torch.Tensor)
+            and not isinstance(image, PIL.Image.Image)
+            and not isinstance(image, list)
+        ):
+            raise ValueError(
+                "`image` has to be of type `torch.FloatTensor` or `PIL.Image.Image` or `List[PIL.Image.Image]` but is"
+                f" {type(image)}"
+            )
+
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+    def prepare_latents(
+        self,
+        batch_size,
+        num_frames,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        shape = (
+            batch_size,
+            num_frames,
+            num_channels_latents // 2,
+            height // self.vae_scale_factor,
+            width // self.vae_scale_factor,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+    # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+    # corresponds to doing no classifier free guidance.
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: Union[PIL.Image.Image, List[PIL.Image.Image], torch.FloatTensor],
+        height: int = 576,
+        width: int = 1024,
+        num_frames: Optional[int] = None,
+        num_inference_steps: int = 25,
+        min_guidance_scale: float = 1.0,
+        max_guidance_scale: float = 3.0,
+        fps: int = 7,
+        motion_bucket_id: int = 127,
+        noise_aug_strength: int = 0.02,
+        decode_chunk_size: Optional[int] = None,
+        num_videos_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        cache_interval: Optional[int] = 1, 
+        cache_branch: Optional[int] = None,
+        return_dict: bool = True,
+    ):
+        r"""
+        The call function to the pipeline for generation.
+
+        Args:
+            image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.FloatTensor`):
+                Image or images to guide image generation. If you provide a tensor, it needs to be compatible with
+                [`CLIPImageProcessor`](https://huggingface.co/lambdalabs/sd-image-variations-diffusers/blob/main/feature_extractor/preprocessor_config.json).
+            height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The width in pixels of the generated image.
+            num_frames (`int`, *optional*):
+                The number of video frames to generate. Defaults to 14 for `stable-video-diffusion-img2vid` and to 25 for `stable-video-diffusion-img2vid-xt`
+            num_inference_steps (`int`, *optional*, defaults to 25):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference. This parameter is modulated by `strength`.
+            min_guidance_scale (`float`, *optional*, defaults to 1.0):
+                The minimum guidance scale. Used for the classifier free guidance with first frame.
+            max_guidance_scale (`float`, *optional*, defaults to 3.0):
+                The maximum guidance scale. Used for the classifier free guidance with last frame.
+            fps (`int`, *optional*, defaults to 7):
+                Frames per second. The rate at which the generated images shall be exported to a video after generation.
+                Note that Stable Diffusion Video's UNet was micro-conditioned on fps-1 during training.
+            motion_bucket_id (`int`, *optional*, defaults to 127):
+                The motion bucket ID. Used as conditioning for the generation. The higher the number the more motion will be in the video.
+            noise_aug_strength (`int`, *optional*, defaults to 0.02):
+                The amount of noise added to the init image, the higher it is the less the video will look like the init image. Increase it for more motion.
+            decode_chunk_size (`int`, *optional*):
+                The number of frames to decode at a time. The higher the chunk size, the higher the temporal consistency
+                between frames, but also the higher the memory consumption. By default, the decoder will decode all frames at once
+                for maximal quality. Reduce `decode_chunk_size` to reduce memory usage.
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor is generated by sampling using the supplied random `generator`.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
+                plain tuple.
+
+        Returns:
+            [`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] is returned,
+                otherwise a `tuple` is returned where the first element is a list of list with the generated frames.
+
+        Examples:
+
+        ```py
+        from diffusers import StableVideoDiffusionPipeline
+        from diffusers.utils import load_image, export_to_video
+
+        pipe = StableVideoDiffusionPipeline.from_pretrained("stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16")
+        pipe.to("cuda")
+
+        image = load_image("https://lh3.googleusercontent.com/y-iFOHfLTwkuQSUegpwDdgKmOjRSTvPxat63dQLB25xkTs4lhIbRUFeNBWZzYf370g=s1200")
+        image = image.resize((1024, 576))
+
+        frames = pipe(image, num_frames=25, decode_chunk_size=8).frames[0]
+        export_to_video(frames, "generated.mp4", fps=7)
+        ```
+        """
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        num_frames = num_frames if num_frames is not None else self.unet.config.num_frames
+        decode_chunk_size = decode_chunk_size if decode_chunk_size is not None else num_frames
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(image, height, width)
+
+        # 2. Define call parameters
+        if isinstance(image, PIL.Image.Image):
+            batch_size = 1
+        elif isinstance(image, list):
+            batch_size = len(image)
+        else:
+            batch_size = image.shape[0]
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = max_guidance_scale > 1.0
+
+        # 3. Encode input image
+        image_embeddings = self._encode_image(image, device, num_videos_per_prompt, do_classifier_free_guidance)
+
+        # NOTE: Stable Diffusion Video was conditioned on fps - 1, which
+        # is why it is reduced here.
+        # See: https://github.com/Stability-AI/generative-models/blob/ed0997173f98eaf8f4edf7ba5fe8f15c6b877fd3/scripts/sampling/simple_video_sample.py#L188
+        fps = fps - 1
+
+        # 4. Encode input image using VAE
+        image = self.image_processor.preprocess(image, height=height, width=width)
+        noise = randn_tensor(image.shape, generator=generator, device=image.device, dtype=image.dtype)
+        image = image + noise_aug_strength * noise
+
+        needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
+        if needs_upcasting:
+            self.vae.to(dtype=torch.float32)
+
+        image_latents = self._encode_vae_image(image, device, num_videos_per_prompt, do_classifier_free_guidance)
+        image_latents = image_latents.to(image_embeddings.dtype)
+
+        # cast back to fp16 if needed
+        if needs_upcasting:
+            self.vae.to(dtype=torch.float16)
+
+        # Repeat the image latents for each frame so we can concatenate them with the noise
+        # image_latents [batch, channels, height, width] ->[batch, num_frames, channels, height, width]
+        image_latents = image_latents.unsqueeze(1).repeat(1, num_frames, 1, 1, 1)
+
+        # 5. Get Added Time IDs
+        added_time_ids = self._get_add_time_ids(
+            fps,
+            motion_bucket_id,
+            noise_aug_strength,
+            image_embeddings.dtype,
+            batch_size,
+            num_videos_per_prompt,
+            do_classifier_free_guidance,
+        )
+        added_time_ids = added_time_ids.to(device)
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_videos_per_prompt,
+            num_frames,
+            num_channels_latents,
+            height,
+            width,
+            image_embeddings.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 7. Prepare guidance scale
+        guidance_scale = torch.linspace(min_guidance_scale, max_guidance_scale, num_frames).unsqueeze(0)
+        guidance_scale = guidance_scale.to(device, latents.dtype)
+        guidance_scale = guidance_scale.repeat(batch_size * num_videos_per_prompt, 1)
+        guidance_scale = _append_dims(guidance_scale, latents.ndim)
+
+        self._guidance_scale = guidance_scale
+
+        cache_features = None
+        interval_seq = list(range(0, num_inference_steps, cache_interval))
+        interval_seq = sorted(interval_seq)
+
+        # 8. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                # Concatenate image_latents over channels dimention
+                latent_model_input = torch.cat([latent_model_input, image_latents], dim=2)
+
+                if i in interval_seq:
+                    cache_features = None
+
+                # predict the noise residual
+                noise_pred, cache_features = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=image_embeddings,
+                    added_time_ids=added_time_ids,
+                    cache_features=cache_features,
+                    cache_branch=cache_branch,
+                    return_dict=False,
+                )
+
+                # perform guidance
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+        if not output_type == "latent":
+            # cast back to fp16 if needed
+            if needs_upcasting:
+                self.vae.to(dtype=torch.float16)
+            frames = self.decode_latents(latents, num_frames, decode_chunk_size)
+            frames = tensor2vid(frames, self.image_processor, output_type=output_type)
+        else:
+            frames = latents
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return frames
+
+        return StableVideoDiffusionPipelineOutput(frames=frames)
+
+
+# resizing utils
+# TODO: clean up later
+def _resize_with_antialiasing(input, size, interpolation="bicubic", align_corners=True):
+    h, w = input.shape[-2:]
+    factors = (h / size[0], w / size[1])
+
+    # First, we have to determine sigma
+    # Taken from skimage: https://github.com/scikit-image/scikit-image/blob/v0.19.2/skimage/transform/_warps.py#L171
+    sigmas = (
+        max((factors[0] - 1.0) / 2.0, 0.001),
+        max((factors[1] - 1.0) / 2.0, 0.001),
+    )
+
+    # Now kernel size. Good results are for 3 sigma, but that is kind of slow. Pillow uses 1 sigma
+    # https://github.com/python-pillow/Pillow/blob/master/src/libImaging/Resample.c#L206
+    # But they do it in the 2 passes, which gives better results. Let's try 2 sigmas for now
+    ks = int(max(2.0 * 2 * sigmas[0], 3)), int(max(2.0 * 2 * sigmas[1], 3))
+
+    # Make sure it is odd
+    if (ks[0] % 2) == 0:
+        ks = ks[0] + 1, ks[1]
+
+    if (ks[1] % 2) == 0:
+        ks = ks[0], ks[1] + 1
+
+    input = _gaussian_blur2d(input, ks, sigmas)
+
+    output = torch.nn.functional.interpolate(input, size=size, mode=interpolation, align_corners=align_corners)
+    return output
+
+
+def _compute_padding(kernel_size):
+    """Compute padding tuple."""
+    # 4 or 6 ints:  (padding_left, padding_right,padding_top,padding_bottom)
+    # https://pytorch.org/docs/stable/nn.html#torch.nn.functional.pad
+    if len(kernel_size) < 2:
+        raise AssertionError(kernel_size)
+    computed = [k - 1 for k in kernel_size]
+
+    # for even kernels we need to do asymmetric padding :(
+    out_padding = 2 * len(kernel_size) * [0]
+
+    for i in range(len(kernel_size)):
+        computed_tmp = computed[-(i + 1)]
+
+        pad_front = computed_tmp // 2
+        pad_rear = computed_tmp - pad_front
+
+        out_padding[2 * i + 0] = pad_front
+        out_padding[2 * i + 1] = pad_rear
+
+    return out_padding
+
+
+def _filter2d(input, kernel):
+    # prepare kernel
+    b, c, h, w = input.shape
+    tmp_kernel = kernel[:, None, ...].to(device=input.device, dtype=input.dtype)
+
+    tmp_kernel = tmp_kernel.expand(-1, c, -1, -1)
+
+    height, width = tmp_kernel.shape[-2:]
+
+    padding_shape: list[int] = _compute_padding([height, width])
+    input = torch.nn.functional.pad(input, padding_shape, mode="reflect")
+
+    # kernel and input tensor reshape to align element-wise or batch-wise params
+    tmp_kernel = tmp_kernel.reshape(-1, 1, height, width)
+    input = input.view(-1, tmp_kernel.size(0), input.size(-2), input.size(-1))
+
+    # convolve the tensor with the kernel.
+    output = torch.nn.functional.conv2d(input, tmp_kernel, groups=tmp_kernel.size(0), padding=0, stride=1)
+
+    out = output.view(b, c, h, w)
+    return out
+
+
+def _gaussian(window_size: int, sigma):
+    if isinstance(sigma, float):
+        sigma = torch.tensor([[sigma]])
+
+    batch_size = sigma.shape[0]
+
+    x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) - window_size // 2).expand(batch_size, -1)
+
+    if window_size % 2 == 0:
+        x = x + 0.5
+
+    gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0)))
+
+    return gauss / gauss.sum(-1, keepdim=True)
+
+
+def _gaussian_blur2d(input, kernel_size, sigma):
+    if isinstance(sigma, tuple):
+        sigma = torch.tensor([sigma], dtype=input.dtype)
+    else:
+        sigma = sigma.to(dtype=input.dtype)
+
+    ky, kx = int(kernel_size[0]), int(kernel_size[1])
+    bs = sigma.shape[0]
+    kernel_x = _gaussian(kx, sigma[:, 1].view(bs, 1))
+    kernel_y = _gaussian(ky, sigma[:, 0].view(bs, 1))
+    out_x = _filter2d(input, kernel_x[..., None, :])
+    out = _filter2d(out_x, kernel_y[..., None])
+
+    return out

ixformer_sdk/contrib/DeepCache/svd/unet_spatio_temporal_condition.py

@@ -0,0 +1,566 @@
+from dataclasses import dataclass
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import UNet2DConditionLoadersMixin
+from diffusers.utils import BaseOutput, logging
+from diffusers.models.attention_processor import CROSS_ATTENTION_PROCESSORS, AttentionProcessor, AttnProcessor
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_utils import ModelMixin
+
+from .unet_3d_blocks import UNetMidBlockSpatioTemporal, get_down_block, get_up_block
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class UNetSpatioTemporalConditionOutput(BaseOutput):
+    """
+    The output of [`UNetSpatioTemporalConditionModel`].
+
+    Args:
+        sample (`torch.FloatTensor` of shape `(batch_size, num_frames, num_channels, height, width)`):
+            The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.
+    """
+
+    sample: torch.FloatTensor = None
+
+
+class UNetSpatioTemporalConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):
+    r"""
+    A conditional Spatio-Temporal UNet model that takes a noisy video frames, conditional state, and a timestep and returns a sample
+    shaped output.
+
+    This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+    for all models (such as downloading or saving).
+
+    Parameters:
+        sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
+            Height and width of input/output sample.
+        in_channels (`int`, *optional*, defaults to 8): Number of channels in the input sample.
+        out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
+        down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "CrossAttnDownBlockSpatioTemporal", "DownBlockSpatioTemporal")`):
+            The tuple of downsample blocks to use.
+        up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal")`):
+            The tuple of upsample blocks to use.
+        block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
+            The tuple of output channels for each block.
+        addition_time_embed_dim: (`int`, defaults to 256):
+            Dimension to to encode the additional time ids.
+        projection_class_embeddings_input_dim (`int`, defaults to 768):
+            The dimension of the projection of encoded `added_time_ids`.
+        layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
+        cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
+            The dimension of the cross attention features.
+        transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1):
+            The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
+            [`~models.unet_3d_blocks.CrossAttnDownBlockSpatioTemporal`], [`~models.unet_3d_blocks.CrossAttnUpBlockSpatioTemporal`],
+            [`~models.unet_3d_blocks.UNetMidBlockSpatioTemporal`].
+        num_attention_heads (`int`, `Tuple[int]`, defaults to `(5, 10, 10, 20)`):
+            The number of attention heads.
+        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        sample_size: Optional[int] = None,
+        in_channels: int = 8,
+        out_channels: int = 4,
+        down_block_types: Tuple[str] = (
+            "CrossAttnDownBlockSpatioTemporal",
+            "CrossAttnDownBlockSpatioTemporal",
+            "CrossAttnDownBlockSpatioTemporal",
+            "DownBlockSpatioTemporal",
+        ),
+        up_block_types: Tuple[str] = (
+            "UpBlockSpatioTemporal",
+            "CrossAttnUpBlockSpatioTemporal",
+            "CrossAttnUpBlockSpatioTemporal",
+            "CrossAttnUpBlockSpatioTemporal",
+        ),
+        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
+        addition_time_embed_dim: int = 256,
+        projection_class_embeddings_input_dim: int = 768,
+        layers_per_block: Union[int, Tuple[int]] = 2,
+        cross_attention_dim: Union[int, Tuple[int]] = 1024,
+        transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1,
+        num_attention_heads: Union[int, Tuple[int]] = (5, 10, 10, 20),
+        num_frames: int = 25,
+    ):
+        super().__init__()
+
+        self.sample_size = sample_size
+
+        # Check inputs
+        if len(down_block_types) != len(up_block_types):
+            raise ValueError(
+                f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
+            )
+
+        if len(block_out_channels) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
+            )
+
+        if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
+            )
+
+        if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):
+            raise ValueError(
+                f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}."
+            )
+
+        # input
+        self.conv_in = nn.Conv2d(
+            in_channels,
+            block_out_channels[0],
+            kernel_size=3,
+            padding=1,
+        )
+
+        # time
+        time_embed_dim = block_out_channels[0] * 4
+
+        self.time_proj = Timesteps(block_out_channels[0], True, downscale_freq_shift=0)
+        timestep_input_dim = block_out_channels[0]
+
+        self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
+
+        self.add_time_proj = Timesteps(addition_time_embed_dim, True, downscale_freq_shift=0)
+        self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+
+        self.down_blocks = nn.ModuleList([])
+        self.up_blocks = nn.ModuleList([])
+
+        if isinstance(num_attention_heads, int):
+            num_attention_heads = (num_attention_heads,) * len(down_block_types)
+
+        if isinstance(cross_attention_dim, int):
+            cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
+
+        if isinstance(layers_per_block, int):
+            layers_per_block = [layers_per_block] * len(down_block_types)
+
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
+
+        blocks_time_embed_dim = time_embed_dim
+
+        # down
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+
+            down_block = get_down_block(
+                down_block_type,
+                num_layers=layers_per_block[i],
+                transformer_layers_per_block=transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                temb_channels=blocks_time_embed_dim,
+                add_downsample=not is_final_block,
+                resnet_eps=1e-5,
+                cross_attention_dim=cross_attention_dim[i],
+                num_attention_heads=num_attention_heads[i],
+                resnet_act_fn="silu",
+            )
+            self.down_blocks.append(down_block)
+
+        # mid
+        self.mid_block = UNetMidBlockSpatioTemporal(
+            block_out_channels[-1],
+            temb_channels=blocks_time_embed_dim,
+            transformer_layers_per_block=transformer_layers_per_block[-1],
+            cross_attention_dim=cross_attention_dim[-1],
+            num_attention_heads=num_attention_heads[-1],
+        )
+
+        # count how many layers upsample the images
+        self.num_upsamplers = 0
+
+        # up
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        reversed_num_attention_heads = list(reversed(num_attention_heads))
+        reversed_layers_per_block = list(reversed(layers_per_block))
+        reversed_cross_attention_dim = list(reversed(cross_attention_dim))
+        reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block))
+
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            is_final_block = i == len(block_out_channels) - 1
+
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
+
+            # add upsample block for all BUT final layer
+            if not is_final_block:
+                add_upsample = True
+                self.num_upsamplers += 1
+            else:
+                add_upsample = False
+
+            up_block = get_up_block(
+                up_block_type,
+                num_layers=reversed_layers_per_block[i] + 1,
+                transformer_layers_per_block=reversed_transformer_layers_per_block[i],
+                in_channels=input_channel,
+                out_channels=output_channel,
+                prev_output_channel=prev_output_channel,
+                temb_channels=blocks_time_embed_dim,
+                add_upsample=add_upsample,
+                resnet_eps=1e-5,
+                resolution_idx=i,
+                cross_attention_dim=reversed_cross_attention_dim[i],
+                num_attention_heads=reversed_num_attention_heads[i],
+                resnet_act_fn="silu",
+            )
+            self.up_blocks.append(up_block)
+            prev_output_channel = output_channel
+
+        # out
+        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=32, eps=1e-5)
+        self.conv_act = nn.SiLU()
+
+        self.conv_out = nn.Conv2d(
+            block_out_channels[0],
+            out_channels,
+            kernel_size=3,
+            padding=1,
+        )
+
+    @property
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(
+            name: str,
+            module: torch.nn.Module,
+            processors: Dict[str, AttentionProcessor],
+        ):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True)
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    def set_attn_processor(self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(proc.__class__ in CROSS_ATTENTION_PROCESSORS for proc in self.attn_processors.values()):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    # Copied from diffusers.models.unet_3d_condition.UNet3DConditionModel.enable_forward_chunking
+    def enable_forward_chunking(self, chunk_size: Optional[int] = None, dim: int = 0) -> None:
+        """
+        Sets the attention processor to use [feed forward
+        chunking](https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers).
+
+        Parameters:
+            chunk_size (`int`, *optional*):
+                The chunk size of the feed-forward layers. If not specified, will run feed-forward layer individually
+                over each tensor of dim=`dim`.
+            dim (`int`, *optional*, defaults to `0`):
+                The dimension over which the feed-forward computation should be chunked. Choose between dim=0 (batch)
+                or dim=1 (sequence length).
+        """
+        if dim not in [0, 1]:
+            raise ValueError(f"Make sure to set `dim` to either 0 or 1, not {dim}")
+
+        # By default chunk size is 1
+        chunk_size = chunk_size or 1
+
+        def fn_recursive_feed_forward(module: torch.nn.Module, chunk_size: int, dim: int):
+            if hasattr(module, "set_chunk_feed_forward"):
+                module.set_chunk_feed_forward(chunk_size=chunk_size, dim=dim)
+
+            for child in module.children():
+                fn_recursive_feed_forward(child, chunk_size, dim)
+
+        for module in self.children():
+            fn_recursive_feed_forward(module, chunk_size, dim)
+
+    def forward(
+        self,
+        sample: torch.FloatTensor,
+        timestep: Union[torch.Tensor, float, int],
+        encoder_hidden_states: torch.Tensor,
+        added_time_ids: torch.Tensor,
+        cache_features: Optional[torch.Tensor] = None,
+        cache_branch: Optional[int] = None,
+        return_dict: bool = True,
+    ) -> Union[UNetSpatioTemporalConditionOutput, Tuple]:
+        r"""
+        The [`UNetSpatioTemporalConditionModel`] forward method.
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The noisy input tensor with the following shape `(batch, num_frames, channel, height, width)`.
+            timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input.
+            encoder_hidden_states (`torch.FloatTensor`):
+                The encoder hidden states with shape `(batch, sequence_length, cross_attention_dim)`.
+            added_time_ids: (`torch.FloatTensor`):
+                The additional time ids with shape `(batch, num_additional_ids)`. These are encoded with sinusoidal
+                embeddings and added to the time embeddings.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] instead of a plain
+                tuple.
+        Returns:
+            [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] or `tuple`:
+                If `return_dict` is True, an [`~models.unet_slatio_temporal.UNetSpatioTemporalConditionOutput`] is returned, otherwise
+                a `tuple` is returned where the first element is the sample tensor.
+        """
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        batch_size, num_frames = sample.shape[:2]
+        timesteps = timesteps.expand(batch_size)
+
+        t_emb = self.time_proj(timesteps)
+
+        # `Timesteps` does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=sample.dtype)
+
+        emb = self.time_embedding(t_emb)
+
+        time_embeds = self.add_time_proj(added_time_ids.flatten())
+        time_embeds = time_embeds.reshape((batch_size, -1))
+        time_embeds = time_embeds.to(emb.dtype)
+        aug_emb = self.add_embedding(time_embeds)
+        emb = emb + aug_emb
+
+        # Flatten the batch and frames dimensions
+        # sample: [batch, frames, channels, height, width] -> [batch * frames, channels, height, width]
+        sample = sample.flatten(0, 1)
+        # Repeat the embeddings num_video_frames times
+        # emb: [batch, channels] -> [batch * frames, channels]
+        emb = emb.repeat_interleave(num_frames, dim=0)
+        # encoder_hidden_states: [batch, 1, channels] -> [batch * frames, 1, channels]
+        encoder_hidden_states = encoder_hidden_states.repeat_interleave(num_frames, dim=0)
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
+        image_only_indicator = torch.zeros(batch_size, num_frames, dtype=sample.dtype, device=sample.device)
+
+        # Branch: 4 down_blocks, each with 3 skip connections. Here we ignore the first skip branch, whose computations only has up_blocks but without down_blocks.
+        if cache_branch is not None:
+            each_module_num = len(self.down_blocks[0].resnets) + 1
+            down_cache_block_idx = cache_branch // each_module_num
+            down_cache_module_idx = cache_branch % each_module_num
+
+            up_cache_block_idx = len(self.up_blocks) - 1 - down_cache_block_idx
+            up_cache_module_idx = 1 - down_cache_module_idx
+            if down_cache_module_idx == each_module_num - 1:
+                up_cache_block_idx -= 1
+                up_cache_module_idx = 2
+            
+        if cache_features is not None:
+            # 3. down
+            down_block_res_samples = (sample,)
+            for block_id, downsample_block in enumerate(self.down_blocks):
+                if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                    sample, res_samples = downsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        encoder_hidden_states=encoder_hidden_states,
+                        image_only_indicator=image_only_indicator,
+                        exist_module_idx=down_cache_module_idx if down_cache_block_idx == block_id else None
+                    )
+                else:
+                    sample, res_samples = downsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        image_only_indicator=image_only_indicator,
+                        exist_module_idx=down_cache_module_idx if down_cache_block_idx == block_id else None
+                    )
+                    
+                down_block_res_samples += res_samples
+                if down_cache_block_idx == block_id:
+                    break
+
+            # 4. no mid
+            sample = cache_features
+
+            # 5. up
+            for i, upsample_block in enumerate(self.up_blocks):
+                if i < up_cache_block_idx:
+                    continue
+                
+                if i == up_cache_block_idx:
+                    trunc_res_samples_len = len(upsample_block.resnets) - up_cache_module_idx
+                else:
+                    trunc_res_samples_len = len(upsample_block.resnets)
+                
+                res_samples = down_block_res_samples[-trunc_res_samples_len :]
+                down_block_res_samples = down_block_res_samples[: -trunc_res_samples_len]
+
+                if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                    sample, _ = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        encoder_hidden_states=encoder_hidden_states,
+                        image_only_indicator=image_only_indicator,
+                        enter_module_idx=up_cache_module_idx if i == up_cache_block_idx else None
+                    )
+                else:
+                    sample, _ = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        image_only_indicator=image_only_indicator,
+                        enter_module_idx=up_cache_module_idx if i == up_cache_block_idx else None
+                    )
+        else:
+            # 3. down
+            down_block_res_samples = (sample,)
+            for downsample_block in self.down_blocks:
+                if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                    sample, res_samples = downsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        encoder_hidden_states=encoder_hidden_states,
+                        image_only_indicator=image_only_indicator,
+                    )
+                else:
+                    sample, res_samples = downsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        image_only_indicator=image_only_indicator,
+                    )
+
+                down_block_res_samples += res_samples
+
+            # 4. mid
+            sample = self.mid_block(
+                hidden_states=sample,
+                temb=emb,
+                encoder_hidden_states=encoder_hidden_states,
+                image_only_indicator=image_only_indicator,
+            )
+
+            # 5. up
+            for i, upsample_block in enumerate(self.up_blocks):
+                res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
+                down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
+
+                if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                    sample, current_record_f = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        encoder_hidden_states=encoder_hidden_states,
+                        image_only_indicator=image_only_indicator,
+                    )
+                else:
+                    sample, current_record_f = upsample_block(
+                        hidden_states=sample,
+                        temb=emb,
+                        res_hidden_states_tuple=res_samples,
+                        image_only_indicator=image_only_indicator,
+                    )
+
+                if cache_branch is not None and i == up_cache_block_idx:
+                    cache_features = current_record_f[up_cache_module_idx]
+
+        # 6. post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        # 7. Reshape back to original shape
+        sample = sample.reshape(batch_size, num_frames, *sample.shape[1:])
+
+        if not return_dict:
+            return (sample, cache_features)
+
+        return UNetSpatioTemporalConditionOutput(sample=sample)

ixformer_sdk/contrib/comfy/unet_model_wrapper.py

@@ -0,0 +1,496 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import ixformer.functions as ixf_F
+
+
+def time_embed(t_emb, weight1, bias1, weight2, bias2):
+    # unet time_emd
+    # linear + silu + linear
+    emb = ixf_F.act_bias_mm(
+        t_emb, weight1, act_type="silu", bias=bias1, scale=1, trans_format="TN"
+    )
+    emb = ixf_F.act_bias_mm(
+        emb, weight2, act_type="none", bias=bias2, scale=1, trans_format="TN"
+    )
+    return emb
+
+
+def ixf_layer_norm(input, normalized_shape, weight=None, bias=None, eps=1e-05):
+    return ixf_F.layernorm(input, weight, bias, normalized_shape)
+
+
+def ixf_pt_scaled_dot_product_attention(
+    query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False
+):
+    if (
+        not query.is_contiguous()
+        and query.transpose(1, 2).is_contiguous()
+        and key.transpose(1, 2).is_contiguous()
+        and value.transpose(1, 2).is_contiguous()
+        and attn_mask is None
+    ):
+
+        batch_size, head_num, seq_len_q, head_dim = query.shape
+        _, _, seq_len_k, _ = key.shape
+
+        query = query.transpose(1, 2).view(batch_size * seq_len_q, head_num, head_dim)
+        key = key.transpose(1, 2).view(batch_size * seq_len_k, head_num, head_dim)
+        value = value.transpose(1, 2).view(batch_size * seq_len_k, head_num, head_dim)
+
+        cu_seqlens_q = torch.arange(
+            0,
+            seq_len_q * (batch_size + 1),
+            seq_len_q,
+            dtype=torch.int32,
+            device=query.device,
+        )
+        if seq_len_q == seq_len_k:
+            cu_seqlens_k = cu_seqlens_q
+        else:
+            cu_seqlens_k = torch.arange(
+                0,
+                seq_len_k * (batch_size + 1),
+                seq_len_k,
+                dtype=torch.int32,
+                device=query.device,
+            )
+
+        res = ixf_F.flash_attn_varlen_func(
+            query,
+            key,
+            value,
+            cu_seqlens_q.int(),
+            cu_seqlens_k.int(),
+            seq_len_q,
+            seq_len_k,
+        )
+        res = res.view(batch_size, seq_len_q, head_num, head_dim).transpose(1, 2)
+        return res
+
+    if not query.is_contiguous():
+        query = query.contiguous()
+    if not key.is_contiguous():
+        key = key.contiguous()
+    if not value.is_contiguous():
+        value = value.contiguous()
+    return ixf_F.scaled_dot_product_attention(
+        query, key, value, attn_mask=attn_mask, is_causal=is_causal
+    )
+
+
+class UnetIxformerFunction:
+    def __init__(self) -> None:
+        self.ixf_linear = ixf_F.linear
+        self.pt_linear = F.linear
+        self.pt_layer_norm = F.layer_norm
+        self.pt_scaled_dot_product_attention = F.scaled_dot_product_attention
+
+    def __enter__(self):
+        F.linear = self.ixf_linear
+        F.layer_norm = ixf_layer_norm
+        F.scaled_dot_product_attention = ixf_pt_scaled_dot_product_attention
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        F.linear = self.pt_linear
+        F.layer_norm = self.pt_layer_norm
+        F.scaled_dot_product_attention = self.pt_scaled_dot_product_attention
+        if exc_tb is not None:
+            print(f"{exc_type} {exc_val}")
+            return False
+        return True
+
+
+def ForwardWrapper(fun):
+    def wrap(*args, **kwargs):
+        with UnetIxformerFunction() as w:
+            return fun(*args, **kwargs)
+
+    return wrap
+
+
+class IxformerComfyWrapper(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.is_ixf_wrapper = True
+
+
+class Conv2dNhwcWrapper(IxformerComfyWrapper):
+    def __init__(self, module):
+        super().__init__()
+        module.weight.data = module.weight.permute(0, 2, 3, 1).contiguous()
+        module.bias.data = module.bias.float()
+        self.weight = module.weight.data
+        self.bias = module.bias.data
+        self.stride = module.stride
+        self.padding = module.padding
+        self.dilation = module.dilation
+        self.groups = module.groups
+
+    def forward(self, x):
+        h2 = ixf_F.conv2d(
+            x,
+            self.weight,
+            self.bias,
+            self.stride,
+            self.padding,
+            self.dilation,
+            self.groups,
+        )
+        return h2
+
+
+class ResBlockNhwcWrapper(IxformerComfyWrapper):
+    def __init__(self, module) -> None:
+        super().__init__()
+        assert not module.updown
+        assert not module.use_scale_shift_norm
+        assert not module.skip_t_emb
+        assert not module.exchange_temb_dims
+
+        if isinstance(module.skip_connection, nn.Identity):
+            self.skip_connection = module.skip_connection
+        elif get_class_name(module.skip_connection) == "Conv2d":
+            self.skip_connection = Conv2dNhwcWrapper(module.skip_connection)
+        else:
+            raise NotImplementedError(
+                f"ResBlockNhwcWrapper support Conv2d or nn.Identity, but got {module.skip_connection}"
+            )
+
+        self.in_layers = module.in_layers
+        self.out_layers = module.out_layers
+        self.emb_layers = module.emb_layers
+        self.in_layers_conv = Conv2dNhwcWrapper(module.in_layers[2])
+        self.out_layers_conv = Conv2dNhwcWrapper(module.out_layers[3])
+
+    def forward(self, x, emb):
+        # x: nhwc
+        x1 = x
+        # print(x1.shape)
+        # fused group_norm silu
+        h = ixf_F.group_norm(
+            x1,  # nchw->nhwc
+            self.in_layers[0].num_groups,
+            self.in_layers[0].weight,
+            self.in_layers[0].bias,
+            format=False,
+            act_type=1,
+        )
+        h = self.in_layers_conv(h)
+
+        emb_out = self.emb_layers(emb)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+
+        h = h + emb_out.permute(0, 2, 3, 1)
+        # print(h.shape)
+        h = ixf_F.group_norm(
+            h,
+            self.out_layers[0].num_groups,
+            self.out_layers[0].weight,
+            self.out_layers[0].bias,
+            format=False,
+            act_type=1,
+        )
+
+        h = self.out_layers[2](h)
+        h = self.out_layers_conv(h)
+        # TODO: support other skip_connection
+        return self.skip_connection(x) + h
+
+
+class DownsampleNhwcWrapper(IxformerComfyWrapper):
+    def __init__(self, module) -> None:
+        # TODO: support avg_pool_nd
+        super().__init__()
+        assert module.use_conv
+        self.channels = module.channels
+        self.op = Conv2dNhwcWrapper(module.op)
+
+    def forward(self, x):
+        assert x.shape[-1] == self.channels
+        return self.op(x)
+
+
+def ffn_forward(self, x):
+    if get_class_name(self.net[0]) == "GEGLU":
+        net = self.net[1:]
+        geglu_net = self.net[0]
+        x = geglu_net.proj(x)
+        x = ixf_F.gelu_and_mul(x)
+        return net(x)
+    else:
+        return self.net(x)
+
+
+# ComfyUI/comfy/ldm/modules/attention.py `class BasicTransformerBlock(nn.Module)`
+def transformer_block_forward(self, x, context=None, transformer_options={}):
+    extra_options = {}
+    block = transformer_options.get("block", None)
+    block_index = transformer_options.get("block_index", 0)
+    transformer_patches = {}
+    transformer_patches_replace = {}
+
+    for k in transformer_options:
+        if k == "patches":
+            transformer_patches = transformer_options[k]
+        elif k == "patches_replace":
+            transformer_patches_replace = transformer_options[k]
+        else:
+            extra_options[k] = transformer_options[k]
+
+    extra_options["n_heads"] = self.n_heads
+    extra_options["dim_head"] = self.d_head
+
+    if self.ff_in:
+        x_skip = x
+        x = self.ff_in(self.norm_in(x))
+        if self.is_res:
+            x += x_skip
+
+    n = self.norm1(x)
+    if self.disable_self_attn:
+        context_attn1 = context
+    else:
+        context_attn1 = None
+    value_attn1 = None
+
+    if "attn1_patch" in transformer_patches:
+        patch = transformer_patches["attn1_patch"]
+        if context_attn1 is None:
+            context_attn1 = n
+        value_attn1 = context_attn1
+        for p in patch:
+            n, context_attn1, value_attn1 = p(
+                n, context_attn1, value_attn1, extra_options
+            )
+
+    if block is not None:
+        transformer_block = (block[0], block[1], block_index)
+    else:
+        transformer_block = None
+
+    attn1_replace_patch = transformer_patches_replace.get("attn1", {})
+    block_attn1 = transformer_block
+    if block_attn1 not in attn1_replace_patch:
+        block_attn1 = block
+
+    if block_attn1 in attn1_replace_patch:
+        if context_attn1 is None:
+            context_attn1 = n
+            value_attn1 = n
+        n = self.attn1.to_q(n)
+        context_attn1 = self.attn1.to_k(context_attn1)
+        value_attn1 = self.attn1.to_v(value_attn1)
+        n = attn1_replace_patch[block_attn1](
+            n, context_attn1, value_attn1, extra_options
+        )
+        n = self.attn1.to_out(n)
+    else:
+        n = self.attn1(n, context=context_attn1, value=value_attn1)
+
+    if "attn1_output_patch" in transformer_patches:
+        patch = transformer_patches["attn1_output_patch"]
+        for p in patch:
+            n = p(n, extra_options)
+
+    x += n
+    if "middle_patch" in transformer_patches:
+        patch = transformer_patches["middle_patch"]
+        for p in patch:
+            x = p(x, extra_options)
+
+    if self.attn2 is not None:
+        n = self.norm2(x)
+        if self.switch_temporal_ca_to_sa:
+            context_attn2 = n
+        else:
+            context_attn2 = context
+        value_attn2 = None
+        if "attn2_patch" in transformer_patches:
+            patch = transformer_patches["attn2_patch"]
+            value_attn2 = context_attn2
+            for p in patch:
+                n, context_attn2, value_attn2 = p(
+                    n, context_attn2, value_attn2, extra_options
+                )
+
+        attn2_replace_patch = transformer_patches_replace.get("attn2", {})
+        block_attn2 = transformer_block
+        if block_attn2 not in attn2_replace_patch:
+            block_attn2 = block
+
+        if block_attn2 in attn2_replace_patch:
+            if value_attn2 is None:
+                value_attn2 = context_attn2
+            n = self.attn2.to_q(n)
+            context_attn2 = self.attn2.to_k(context_attn2)
+            value_attn2 = self.attn2.to_v(value_attn2)
+            n = attn2_replace_patch[block_attn2](
+                n, context_attn2, value_attn2, extra_options
+            )
+            n = self.attn2.to_out(n)
+        else:
+            n = self.attn2(n, context=context_attn2, value=value_attn2)
+
+    if "attn2_output_patch" in transformer_patches:
+        patch = transformer_patches["attn2_output_patch"]
+        for p in patch:
+            n = p(n, extra_options)
+
+    # x += n
+    # if self.is_res:
+    #     x_skip = x
+    # x = self.ff(self.norm3(x))
+
+    x, x_skip = ixf_F.residual_layer_norm(
+        n,
+        self.norm3.normalized_shape,
+        self.norm3.weight,
+        self.norm3.bias,
+        x,
+        eps=self.norm3.eps,
+        is_post_ln=False,
+    )
+    x = ffn_forward(self.ff, x)
+
+    # x = ffn_forward(self.ff, self.norm3(x))
+    if self.is_res:
+        x += x_skip
+
+    return x
+
+
+class SpatialTransformerNhwcWrapper(IxformerComfyWrapper):
+    def __init__(self, module):
+        super().__init__()
+        self.use_linear = module.use_linear
+        self.transformer_blocks = module.transformer_blocks
+        self.norm = module.norm
+        if not self.use_linear:
+            self.proj_in = Conv2dNhwcWrapper(module.proj_in)
+            self.proj_out = Conv2dNhwcWrapper(module.proj_out)
+        else:
+            self.proj_in = module.proj_in
+            self.proj_out = module.proj_out
+
+    @ForwardWrapper
+    def forward(self, x, context=None, transformer_options={}):
+        # note: if no context is given, cross-attention defaults to self-attention
+        if not isinstance(context, list):
+            context = [context] * len(self.transformer_blocks)
+
+        b, h, w, c = x.shape
+        x_in = x
+
+        # group_norm
+        x = ixf_F.group_norm(
+            x,
+            self.norm.num_groups,
+            self.norm.weight,
+            self.norm.bias,
+            format=False,
+        )
+        # conv2d
+        if not self.use_linear:
+            x = self.proj_in(x)
+        # n,(hw),c
+        x = x.view(x.shape[0], -1, x.shape[-1])
+        if self.use_linear:
+            x = self.proj_in(x)
+
+        for i, block in enumerate(self.transformer_blocks):
+            transformer_options["block_index"] = i
+            # x = block(x, context=context[i], transformer_options=transformer_options)
+            x = transformer_block_forward(
+                block, x, context=context[i], transformer_options=transformer_options
+            )
+
+        if self.use_linear:
+            x = self.proj_out(x)
+        x = x.view(b, h, w, c)
+        if not self.use_linear:
+            x = self.proj_out(x)
+        return x + x_in
+
+
+class UpsampleNhwcWrapper(IxformerComfyWrapper):
+    def __init__(self, module) -> None:
+        # TODO: support mhwc interpolate
+        super().__init__()
+        self.dims = module.dims
+        self.use_conv = module.use_conv
+        self.channels = module.channels
+        if self.use_conv:
+            self.conv = Conv2dNhwcWrapper(module.conv)
+
+    def forward(self, x, output_shape=None):
+        # print("================== Upsample is running ==================")
+        assert x.shape[-1] == self.channels
+        assert len(x.shape) == 4
+
+        # nhwc -> nchw
+        if output_shape is not None:
+            assert len(output_shape) == 4
+            output_shape = [
+                output_shape[0],
+                output_shape[3],
+                output_shape[1],
+                output_shape[2],
+            ]
+        x = x.permute(0, 3, 1, 2).contiguous()
+        if self.dims == 3:
+            shape = [x.shape[2], x.shape[3] * 2, x.shape[4] * 2]
+            if output_shape is not None:
+                shape[1] = output_shape[3]
+                shape[2] = output_shape[4]
+        else:
+            shape = [x.shape[2] * 2, x.shape[3] * 2]
+            if output_shape is not None:
+                shape[0] = output_shape[2]
+                shape[1] = output_shape[3]
+        # TODO: interpolate 支持 nhwc， 去掉前后转置
+        x = F.interpolate(x, size=shape, mode="nearest")
+        # nchw -> nhwc
+        x = x.permute(0, 2, 3, 1).contiguous()
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+unet_wrappers = {
+    "Conv2d": Conv2dNhwcWrapper,
+    "ResBlock": ResBlockNhwcWrapper,
+    "Downsample": DownsampleNhwcWrapper,
+    "SpatialTransformer": SpatialTransformerNhwcWrapper,
+    "Upsample": UpsampleNhwcWrapper,
+}
+
+
+def get_class_name(module):
+    return module.__class__.__name__
+
+
+def module_wrapper(module):
+    # 将原始的 module 封装为 nhwc 模式
+    module_name = get_class_name(module)
+    assert (
+        module_name == "TimestepEmbedSequential"
+    ), f"ixformer unet_model_wrapper only support 'TimestepEmbedSequential' now, but got {module_name}"
+
+    num_sequential = len(module)
+    for idx_seq in range(num_sequential):
+        sub_module = module[idx_seq]
+        sub_module_name = get_class_name(sub_module)
+        # 判断模块是否已经封装
+        if not getattr(sub_module, "is_ixf_wrapper", False):
+            if sub_module_name in unet_wrappers:
+                module[idx_seq].forward = unet_wrappers[sub_module_name](
+                    sub_module
+                ).forward
+                module[idx_seq].is_ixf_wrapper = True
+            else:
+                raise NotImplementedError(f"{sub_module_name} not support")
+    return module

ixformer_sdk/contrib/flashinfer/__init__.py

@@ -0,0 +1,17 @@
+from .decode import BatchDecodeWithPagedKVCacheWrapper
+from .prefill import (
+    BatchPrefillWithPagedKVCacheWrapper,
+    BatchPrefillWithRaggedKVCacheWrapper,
+)
+
+
+def bmm_fp8():
+    pass
+
+
+def SegmentGEMMWrapper():
+    pass
+
+
+def bmm_fp8():
+    pass

ixformer_sdk/contrib/flashinfer/activation.py

@@ -0,0 +1,29 @@
+import ixformer.inference.functions as ops
+import torch
+
+
+def gelu_and_mul():
+    pass
+
+
+def gelu_tanh_and_mul():
+    pass
+
+
+def silu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
+    r"""Fused SiLU and Mul operation.
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (..., 2 * hidden_size).
+
+    out: Optional[torch.Tensor]
+        The the output tensor, if specified, the kernel will update this tensor inplace.
+
+    Returns
+    -------
+    output: torch.Tensor
+        Output tensor, shape (..., hidden_size).
+    """
+    return ops.silu_and_mul(input=input, output=out)

ixformer_sdk/contrib/flashinfer/cascade.py

@@ -0,0 +1,2 @@
+def merge_state():
+    pass

ixformer_sdk/contrib/flashinfer/decode.py

@@ -0,0 +1,101 @@
+import math
+from typing import Optional, Tuple, Union
+
+import ixformer.inference.functions as ops
+import torch
+
+
+def _grouped_size_compiled_for_decode_kernels(
+    num_qo_heads: int, num_kv_heads: int
+) -> bool:
+    return (num_qo_heads // num_kv_heads) in [1, 2, 4, 8]
+
+
+class BatchDecodeWithPagedKVCacheWrapper:
+    def __init__(
+        self,
+        float_workspace_buffer: torch.Tensor,
+        kv_layout: str = "NHD",
+        use_cuda_graph: bool = False,
+        use_tensor_cores: bool = False,
+    ) -> None:
+        pass
+
+    def plan(
+        self,
+        indptr: torch.Tensor,
+        indices: torch.Tensor,
+        last_page_len: torch.Tensor,
+        num_qo_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        page_size: int,
+        # pos_encoding_mode: str = "NONE",
+        # window_left: int = -1,
+        # logits_soft_cap: Optional[float] = None,
+        data_type: Union[str, torch.dtype] = "float16",
+        q_data_type: Optional[Union[str, torch.dtype]] = None,
+        sm_scale: Optional[float] = None,
+        # rope_scale: Optional[float] = None,
+        # rope_theta: Optional[float] = None,
+        max_seqlen_q: int = None,
+        max_seqlen_k: int = None,
+    ) -> None:
+        self.indptr = indptr
+        self.indices = indices
+        self.last_page_len = last_page_len
+        self.num_qo_heads = num_qo_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+
+        assert page_size == 1
+
+        self.cu_seqlens_q = torch.ones_like(indptr)
+        self.cu_seqlens_q[0] = 0
+        self.cu_seqlens_q = torch.cumsum(self.cu_seqlens_q, dim=0).int()
+
+        self.cu_seqlens_k = indptr
+        if sm_scale is None:
+            sm_scale = 1.0 / math.sqrt(head_dim)
+
+            self.sm_scale = sm_scale
+            self.max_seqlen_q = max_seqlen_q
+            self.max_seqlen_k = max_seqlen_k
+
+    begin_forward = plan
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        paged_kv_cache: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
+        pos_encoding_mode: str = "NONE",
+        q_scale: Optional[float] = None,
+        k_scale: Optional[float] = None,
+        v_scale: Optional[float] = None,
+        window_left: int = -1,
+        logits_soft_cap: Optional[float] = None,
+        sm_scale: Optional[float] = None,
+        rope_scale: Optional[float] = None,
+        rope_theta: Optional[float] = None,
+    ) -> torch.Tensor:
+        k_cache, v_cache = paged_kv_cache
+
+        out = torch.empty_like(q)
+
+        ops.paged_attention_flashinfer(
+            output=out,
+            query=q,
+            paged_kv_data=(k_cache.unsqueeze(1), v_cache.unsqueeze(1)),
+            paged_kv_indptr=self.indptr,
+            paged_kv_indices=self.indices,
+            paged_kv_last_page_len=self.last_page_len,
+            scale=self.sm_scale,
+            max_seq_len=self.max_seqlen_k,
+            kv_cache_format="NHD",
+        )
+
+        return out
+
+    def end_forward(self) -> None:
+        r"""Warning: this function is deprecated and has no effect."""
+        pass

ixformer_sdk/contrib/flashinfer/norm.py

@@ -0,0 +1,61 @@
+import ixformer.inference.functions as ops
+import torch
+
+
+def fused_add_rmsnorm(
+    input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
+):
+    r"""Fused add root mean square normalization.
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (batch_size, hidden_size).
+    residual: torch.Tensor
+        Residual tensor, shape (batch_size, hidden_size).
+    weight: torch.Tensor
+        Weight tensor, shape (hidden_size,).
+    eps: float
+        Epsilon for numerical stability.
+    """
+    return ops.residual_rms_norm(
+        input=input,
+        residual=residual,
+        weight=weight,
+        eps=eps,
+    )
+
+
+def gemma_fused_add_rmsnorm():
+    pass
+
+
+def gemma_rmsnorm():
+    pass
+
+
+def rmsnorm(
+    input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
+) -> torch.Tensor:
+    r"""Root mean square normalization.
+
+    Parameters
+    ----------
+    input: torch.Tensor
+        Input tensor, shape (batch_size, hidden_size).
+    weight: torch.Tensor
+        Weight tensor, shape (hidden_size,).
+    eps: float
+        Epsilon for numerical stability.
+
+    Returns
+    -------
+    output: torch.Tensor
+        Normalized tensor, shape (batch_size, hidden_size).
+    """
+
+    return ops.rms_norm(
+        input=input,
+        weight=weight,
+        eps=eps,
+    )

ixformer_sdk/contrib/flashinfer/prefill.py

@@ -0,0 +1,113 @@
+import math
+from typing import Optional, Tuple, Union
+
+import ixformer._C as ops
+import torch
+
+
+class BatchPrefillWithRaggedKVCacheWrapper:
+    def __init__(
+        self,
+        float_workspace_buffer: torch.Tensor,
+        kv_layout: str = "NHD",
+    ):
+        pass
+
+    def plan(
+        self,
+        qo_indptr: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        num_qo_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        # custom_mask: Optional[torch.Tensor] = None,
+        # packed_custom_mask: Optional[torch.Tensor] = None,
+        causal: bool = True,
+        # pos_encoding_mode: str = "NONE",
+        # allow_fp16_qk_reduction: bool = False,
+        # window_left: int = -1,
+        # logits_soft_cap: Optional[float] = None,
+        sm_scale: Optional[float] = None,
+        # rope_scale: Optional[float] = None,
+        # rope_theta: Optional[float] = None,
+        # q_data_type: str = "float16",
+    ) -> None:
+        batch_size = len(qo_indptr) - 1
+        if len(kv_indptr) != batch_size + 1:
+            raise ValueError(
+                "The kv_indptr length should be equal to qk_indptr length."
+            )
+        self._causal = causal
+        self._sm_scale = sm_scale
+        if sm_scale is None:
+            sm_scale = 1.0 / math.sqrt(head_dim)
+
+        self.cu_seqlens_q = qo_indptr
+        self.cu_seqlens_k = kv_indptr
+        self.num_qo_heads = num_qo_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.max_seqlen_q = max_seqlen_q
+        self.max_seqlen_k = max_seqlen_k
+
+    begin_forward = plan
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        causal: bool = True,
+        # pos_encoding_mode: str = "NONE",
+        # allow_fp16_qk_reduction: bool = False,
+        # window_left: int = -1,
+        logits_soft_cap: Optional[float] = None,
+        sm_scale: Optional[float] = None,
+        # rope_scale: Optional[float] = None,
+        # rope_theta: Optional[float] = None,
+    ) -> torch.Tensor:
+        r"""Warning: This function is deprecated, please use :meth:`run` instead."""
+
+        q = q.view(-1, self.num_qo_heads, self.head_dim)
+        k = k.view(-1, self.num_kv_heads, self.head_dim)
+        v = v.view(-1, self.num_kv_heads, self.head_dim)
+
+        out = torch.empty_like(q)
+
+        assert causal
+        assert (
+            logits_soft_cap is None or logits_soft_cap == 0
+        ), f"logits_soft_cap not supported, but got logits_soft_cap={logits_soft_cap}"
+
+        ops.infer.ixinfer_flash_attn_unpad(
+            q,
+            k,
+            v,
+            out,
+            self.cu_seqlens_q,
+            self.cu_seqlens_k,
+            self.max_seqlen_q,
+            self.max_seqlen_k,
+            causal,
+            False,  # need_lse =False
+            sm_scale,
+            False,
+            None,
+        )
+        return out
+
+    def end_forward(self) -> None:
+        r"""Warning: this function is deprecated and has no effect."""
+        pass
+
+
+class BatchPrefillWithPagedKVCacheWrapper:
+    def __init__(
+        self,
+        float_workspace_buffer: torch.Tensor,
+        kv_layout: str = "NHD",
+        use_cuda_graph: bool = False,
+    ) -> None:
+        pass

ixformer_sdk/contrib/flashinfer/sampling.py

@@ -0,0 +1,14 @@
+def min_p_sampling_from_probs():
+    pass
+
+
+def top_k_renorm_prob():
+    pass
+
+
+def top_k_top_p_sampling_from_probs():
+    pass
+
+
+def top_p_renorm_prob():
+    pass

ixformer_sdk/contrib/tgi/__init__.py

@@ -0,0 +1 @@
+from .fused_moe import fused_moe

ixformer_sdk/contrib/tgi/fused_moe.py

@@ -0,0 +1,430 @@
+import functools
+import json
+import os
+from typing import Any, Dict, Optional, Tuple
+from loguru import logger
+import torch
+import ixformer.inference.functions as ops
+
+CHUNK_SIZE = int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "65536"))
+
+def fused_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], (
+        "Number of tokens mismatch")
+
+    M, _ = hidden_states.shape
+
+    topk_weights = torch.empty(M,
+                               topk,
+                               dtype=torch.float32,
+                               device=hidden_states.device)
+    topk_ids = torch.empty(M,
+                           topk,
+                           dtype=torch.int32,
+                           device=hidden_states.device)
+    token_expert_indicies = torch.empty(M,
+                                        topk,
+                                        dtype=torch.int32,
+                                        device=hidden_states.device)
+    ops.vllm_moe_topk_softmax(
+        topk_weights,
+        topk_ids,
+        token_expert_indicies,
+        gating_output.float(),  # TODO(woosuk): Optimize this.
+    )
+    del token_expert_indicies  # Not used. Will be used in the future.
+
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+    return topk_weights, topk_ids
+
+# This is used by the Deepseek-V2 model
+def grouped_topk(hidden_states: torch.Tensor,
+                 gating_output: torch.Tensor,
+                 topk: int,
+                 renormalize: bool,
+                 num_expert_group: int = 0,
+                 topk_group: int = 0):
+
+    assert hidden_states.shape[0] == gating_output.shape[0], (
+        "Number of tokens mismatch")
+
+    scores = torch.softmax(gating_output, dim=-1)
+    num_token = scores.shape[0]
+    group_scores = scores.view(num_token, num_expert_group,
+                               -1).max(dim=-1).values  # [n, n_group]
+    group_idx = torch.topk(group_scores, k=topk_group, dim=-1,
+                           sorted=False)[1]  # [n, top_k_group]
+    group_mask = torch.zeros_like(group_scores)  # [n, n_group]
+    group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
+    score_mask = group_mask.unsqueeze(-1).expand(
+        num_token, num_expert_group,
+        scores.shape[-1] // num_expert_group).reshape(num_token, -1)  # [n, e]
+    tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0)  # [n, e]
+    topk_weights, topk_ids = torch.topk(tmp_scores,
+                                        k=topk,
+                                        dim=-1,
+                                        sorted=False)
+
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+    return topk_weights, topk_ids
+
+def get_config_file_name(E: int, N: int, dtype: Optional[str]) -> str:
+    device_name = torch.cuda.get_device_name().replace(" ", "_")
+    dtype_selector = "" if not dtype else f",dtype={dtype}"
+    return f"E={E},N={N},device_name={device_name}{dtype_selector}.json"
+
+@functools.lru_cache
+def get_moe_configs(E: int, N: int,
+                    dtype: Optional[str]) -> Optional[Dict[int, Any]]:
+    """
+    Return optimized configurations for the fused MoE kernel.
+
+    The return value will be a dictionary that maps an irregular grid of
+    batch sizes to configurations of the fused_moe kernel. To evaluate the
+    kernel on a given batch size bs, the closest batch size in the grid should
+    be picked and the associated configuration chosen to invoke the kernel.
+    """
+
+    # First look up if an optimized configuration is available in the configs
+    # directory
+    json_file_name = get_config_file_name(E, N, dtype)
+
+    config_file_path = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name)
+    if os.path.exists(config_file_path):
+        with open(config_file_path) as f:
+            logger.info("Using configuration from %s for MoE layer.",
+                        config_file_path)
+            # If a configuration has been found, return it
+            return {int(key): val for key, val in json.load(f).items()}
+
+    # If no optimized configuration is available, we will use the default
+    # configuration
+    return None
+
+def get_default_config(
+    M: int,
+    E: int,
+    N: int,
+    K: int,
+    topk: int,
+    dtype: Optional[str],
+) -> Dict[str, int]:
+    config = {
+        'BLOCK_SIZE_M': 64,
+        'BLOCK_SIZE_N': 64,
+        'BLOCK_SIZE_K': 32,
+        'GROUP_SIZE_M': 8
+    }
+    if M <= E:
+        config = {
+            'BLOCK_SIZE_M': 16,
+            'BLOCK_SIZE_N': 32,
+            'BLOCK_SIZE_K': 64,
+            'GROUP_SIZE_M': 1
+        }
+    numel = M * topk
+    if numel <= 64:
+        config['BLOCK_SIZE_M'] = 32
+    elif numel <= 1024:
+        config['BLOCK_SIZE_M'] = 64
+    else:
+        config['BLOCK_SIZE_M'] = 256
+    return config
+
+def try_get_optimal_moe_config(
+    w1_shape: Tuple[int, ...],
+    w2_shape: Tuple[int, ...],
+    top_k: int,
+    dtype: Optional[str],
+    M: int,
+    override_config: Optional[Dict[str, Any]] = None,
+):
+    if override_config:
+        config = override_config
+    else:
+        # First try to load optimal config from the file
+        E, _, N = w2_shape
+        configs = get_moe_configs(E, N, dtype)
+
+        if configs:
+            # If an optimal configuration map has been found, look up the
+            # optimal config
+            config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
+        else:
+            # Else use the default config
+            config = get_default_config(M, E, N, w1_shape[2], top_k, dtype)
+    return config
+
+def moe_align_block_size(
+        topk_ids: torch.Tensor, block_size: int,
+        num_experts: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Aligns the token distribution across experts to be compatible with block
+    size for matrix multiplication.
+
+    Parameters:
+    - topk_ids: A tensor of shape [total_tokens, top_k] representing the
+        top-k expert indices for each token.
+    - block_size: The block size used in block matrix multiplication.
+    - num_experts: The total number of experts.
+
+    Returns:
+    - sorted_token_ids: A tensor containing the sorted token indices according
+        to their allocated expert.
+    - expert_ids: A tensor indicating the assigned expert index for each block.
+    - num_tokens_post_padded: The total number of tokens after padding,
+        ensuring divisibility by block_size.
+
+    This function pads the number of tokens that each expert needs to process
+    so that it is divisible by block_size.
+    Padding ensures that during block matrix multiplication, the dimensions
+    align correctly.
+
+    Example:
+    Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]],
+    block_size = 4, and num_experts = 4:
+    - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts,
+        with each expert needing to process 3 tokens.
+    - As block_size is 4, we pad 1 token for each expert.
+    - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3].
+    - Then append padding tokens [12, 12, 12, 12] for each block.
+    - After sorting by expert index, we obtain token_ids
+        [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12].
+        Tokens 12 are non-existent (padding) and are ignored in
+        the subsequent matrix multiplication.
+    - The padding ensures that the total number of tokens is now divisible
+        by block_size for proper block matrix operations.
+    """
+    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+    sorted_ids = torch.empty((max_num_tokens_padded, ),
+                             dtype=torch.int32,
+                             device=topk_ids.device)
+    sorted_ids.fill_(topk_ids.numel())
+     # max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
+    max_num_m_blocks = topk_ids.numel() + num_experts
+    expert_ids = torch.empty((max_num_m_blocks, ),
+                             dtype=torch.int32,
+                             device=topk_ids.device)
+    num_tokens_post_pad = torch.empty((1),
+                                      dtype=torch.int32,
+                                      device=topk_ids.device)
+    ops.vllm_moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
+                             expert_ids, num_tokens_post_pad)
+    return sorted_ids, expert_ids, num_tokens_post_pad
+
+def invoke_fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
+                            A_scale: Optional[torch.Tensor],
+                            B_scale: Optional[torch.Tensor],
+                            topk_weights: torch.Tensor, topk_ids: torch.Tensor,
+                            sorted_token_ids: torch.Tensor,
+                            expert_ids: torch.Tensor,
+                            num_tokens_post_padded: torch.Tensor,
+                            mul_routed_weight: bool, top_k: int,
+                            config: Dict[str, Any], compute_type: torch.dtype,
+                            use_fp8: bool) -> None:
+    ops.vllm_invoke_fused_moe_kernel(A, B, C, topk_weights, topk_ids,
+                                     sorted_token_ids,expert_ids, num_tokens_post_padded,
+                                     mul_routed_weight, top_k, config['BLOCK_SIZE_M'])
+
+def fused_experts(hidden_states: torch.Tensor,
+                  w1: torch.Tensor,
+                  w2: torch.Tensor,
+                  topk_weights: torch.Tensor,
+                  topk_ids: torch.Tensor,
+                  inplace: bool = False,
+                  override_config: Optional[Dict[str, Any]] = None,
+                  use_fp8: bool = False,
+                  w1_scale: Optional[torch.Tensor] = None,
+                  w2_scale: Optional[torch.Tensor] = None,
+                  a1_scale: Optional[torch.Tensor] = None,
+                  a2_scale: Optional[torch.Tensor] = None):
+    # Check constraints.
+    assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
+    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
+    assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
+    assert w1.is_contiguous(), "Expert weights1 must be contiguous"
+    assert w2.is_contiguous(), "Expert weights2 must be contiguous"
+    assert hidden_states.dtype in [
+        torch.float32, torch.float16, torch.bfloat16
+    ]
+
+    num_tokens, _ = hidden_states.shape
+    E, N, _ = w1.shape
+    # We execute the fused_moe kernel in chunks to circumvent this issue:
+    # https://github.com/vllm-project/vllm/issues/5938
+    M = min(num_tokens, CHUNK_SIZE)
+
+    get_config_func = functools.partial(
+        try_get_optimal_moe_config,
+        w1.shape,
+        w2.shape,
+        topk_ids.shape[1],
+        "float8" if use_fp8 else None,
+        override_config=override_config,
+    )
+
+    config = get_config_func(M)
+
+    intermediate_cache1 = torch.empty((M, topk_ids.shape[1], N),
+                                      device=hidden_states.device,
+                                      dtype=hidden_states.dtype)
+    intermediate_cache2 = torch.empty((M * topk_ids.shape[1], N // 2),
+                                      device=hidden_states.device,
+                                      dtype=hidden_states.dtype)
+    intermediate_cache3 = torch.empty((M, topk_ids.shape[1], w2.shape[1]),
+                                      device=hidden_states.device,
+                                      dtype=hidden_states.dtype)
+
+    compute_type = (torch.bfloat16
+                    if hidden_states.dtype == torch.bfloat16 else torch.float16)
+
+    if inplace:
+        out_hidden_states = hidden_states
+    else:
+        out_hidden_states = torch.empty_like(hidden_states)
+
+    for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+        begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
+                                          min((chunk + 1) * CHUNK_SIZE,
+                                              num_tokens))
+        curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
+        tokens_in_chunk, _ = curr_hidden_states.shape
+
+        if tokens_in_chunk == 0:
+            break
+
+        if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
+            # Adjust the intermediate cache size and config for the last
+            # chunk. Note that in most cases we only have one chunk
+            # so the cache size and config are already set correctly and
+            # do not need to be adjusted.
+            intermediate_cache1 = intermediate_cache1[:tokens_in_chunk]
+            intermediate_cache2 = intermediate_cache2[:tokens_in_chunk]
+            intermediate_cache3 = intermediate_cache3[:tokens_in_chunk]
+            config = get_config_func(tokens_in_chunk)
+
+        curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+        curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+
+        sorted_token_ids, expert_ids, num_tokens_post_padded = (
+            moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'], E))
+
+        invoke_fused_moe_kernel(curr_hidden_states,
+                                w1,
+                                intermediate_cache1,
+                                a1_scale,
+                                w1_scale,
+                                curr_topk_weights,
+                                curr_topk_ids,
+                                sorted_token_ids,
+                                expert_ids,
+                                num_tokens_post_padded,
+                                False,
+                                topk_ids.shape[1],
+                                config,
+                                compute_type=compute_type,
+                                use_fp8=use_fp8)
+
+        ops.silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
+
+        invoke_fused_moe_kernel(intermediate_cache2,
+                                w2,
+                                intermediate_cache3,
+                                a2_scale,
+                                w2_scale,
+                                curr_topk_weights,
+                                curr_topk_ids,
+                                sorted_token_ids,
+                                expert_ids,
+                                num_tokens_post_padded,
+                                True,
+                                1,
+                                config,
+                                compute_type=compute_type,
+                                use_fp8=use_fp8)
+
+        torch.sum(intermediate_cache3.view(*intermediate_cache3.shape),
+                  dim=1,
+                  out=out_hidden_states[begin_chunk_idx:end_chunk_idx])
+    return out_hidden_states
+
+def fused_moe(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    inplace: bool = False,
+    override_config: Optional[Dict[str, Any]] = None,
+    use_grouped_topk: bool = False,
+    num_expert_group: Optional[int] = None,
+    topk_group: Optional[int] = None,
+    use_fp8: bool = False,
+    w1_scale: Optional[torch.Tensor] = None,
+    w2_scale: Optional[torch.Tensor] = None,
+    a1_scale: Optional[torch.Tensor] = None,
+    a2_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    This function computes a Mixture of Experts (MoE) layer using two sets of
+    weights, w1 and w2, and top-k gating mechanism.
+
+    Parameters:
+    - hidden_states (torch.Tensor): The input tensor to the MoE layer.
+    - w1 (torch.Tensor): The first set of expert weights.
+    - w2 (torch.Tensor): The second set of expert weights.
+    - gating_output (torch.Tensor): The output of the gating operation
+        (before softmax).
+    - topk (int): The number of top-k experts to select.
+    - renormalize (bool): If True, renormalize the top-k weights to sum to 1.
+    - inplace (bool): If True, perform the operation in-place.
+        Defaults to False.
+    - override_config (Optional[Dict[str, Any]]): Optional override
+        for the kernel configuration.
+    - num_expert_group: Optional[int]: additional parameter for grouped_topk
+    - topk_group: Optional[int]: additional parameter for grouped_topk
+    - use_grouped_topk: If True, use grouped_topk instead of fused_topk
+        note: Deepseekv2 model uses grouped_topk
+    - use_fp8 (bool): If True, use fp8 arithmetic to compute the inner
+        products for w1 and w2. Defaults to False.
+    - w1_scale (Optional[torch.Tensor]): Optional scale to be used for
+        w1.
+    - w2_scale (Optional[torch.Tensor]): Optional scale to be used for
+        w2.
+
+    Returns:
+    - torch.Tensor: The output tensor after applying the MoE layer.
+    """
+    # Check constraints.
+    assert gating_output.shape[1] == w1.shape[0], "Number of experts mismatch"
+
+    if use_grouped_topk:
+        assert num_expert_group is not None and topk_group is not None
+        topk_weights, topk_ids = grouped_topk(hidden_states, gating_output,
+                                              topk, renormalize,
+                                              num_expert_group, topk_group)
+    else:
+        topk_weights, topk_ids = fused_topk(hidden_states, gating_output, topk,
+                                            renormalize)
+
+    return fused_experts(hidden_states,
+                         w1,
+                         w2,
+                         topk_weights,
+                         topk_ids,
+                         inplace=inplace,
+                         override_config=override_config,
+                         use_fp8=use_fp8,
+                         w1_scale=w1_scale,
+                         w2_scale=w2_scale,
+                         a1_scale=a1_scale,
+                         a2_scale=a2_scale)

ixformer_sdk/contrib/transformers/__init__.py

@@ -0,0 +1,2 @@
+from .models.bert.modeling_bert import BertForQuestionAnswering
+from .models.t5.modeling_t5 import T5ForConditionalGeneration

ixformer_sdk/contrib/transformers/models/bert/configuration_bert.py

@@ -0,0 +1,150 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  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.
+"""BERT model configuration"""
+
+from collections import OrderedDict
+from typing import Mapping
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.onnx import OnnxConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class BertConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`BertModel`] or a [`TFBertModel`]. It is used to
+    instantiate a BERT model according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the BERT
+    [google-bert/bert-base-uncased](https://huggingface.co/google-bert/bert-base-uncased) 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 30522):
+            Vocabulary size of the BERT model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`BertModel`] or [`TFBertModel`].
+        hidden_size (`int`, *optional*, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        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.
+        intermediate_size (`int`, *optional*, defaults to 3072):
+            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
+        hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"silu"` and `"gelu_new"` are supported.
+        hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
+            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+        attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
+            The dropout ratio for the attention probabilities.
+        max_position_embeddings (`int`, *optional*, defaults to 512):
+            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).
+        type_vocab_size (`int`, *optional*, defaults to 2):
+            The vocabulary size of the `token_type_ids` passed when calling [`BertModel`] or [`TFBertModel`].
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-12):
+            The epsilon used by the layer normalization layers.
+        position_embedding_type (`str`, *optional*, defaults to `"absolute"`):
+            Type of position embedding. Choose one of `"absolute"`, `"relative_key"`, `"relative_key_query"`. For
+            positional embeddings use `"absolute"`. For more information on `"relative_key"`, please refer to
+            [Self-Attention with Relative Position Representations (Shaw et al.)](https://arxiv.org/abs/1803.02155).
+            For more information on `"relative_key_query"`, please refer to *Method 4* in [Improve Transformer Models
+            with Better Relative Position Embeddings (Huang et al.)](https://arxiv.org/abs/2009.13658).
+        is_decoder (`bool`, *optional*, defaults to `False`):
+            Whether the model is used as a decoder or not. If `False`, the model is used as an encoder.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        classifier_dropout (`float`, *optional*):
+            The dropout ratio for the classification head.
+
+    Examples:
+
+    ```python
+    >>> from transformers import BertConfig, BertModel
+
+    >>> # Initializing a BERT google-bert/bert-base-uncased style configuration
+    >>> configuration = BertConfig()
+
+    >>> # Initializing a model (with random weights) from the google-bert/bert-base-uncased style configuration
+    >>> model = BertModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "bert"
+
+    def __init__(
+        self,
+        vocab_size=30522,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.1,
+        max_position_embeddings=512,
+        type_vocab_size=2,
+        initializer_range=0.02,
+        layer_norm_eps=1e-12,
+        pad_token_id=0,
+        position_embedding_type="absolute",
+        use_cache=True,
+        classifier_dropout=None,
+        **kwargs,
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_size = type_vocab_size
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.position_embedding_type = position_embedding_type
+        self.use_cache = use_cache
+        self.classifier_dropout = classifier_dropout
+
+
+class BertOnnxConfig(OnnxConfig):
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        if self.task == "multiple-choice":
+            dynamic_axis = {0: "batch", 1: "choice", 2: "sequence"}
+        else:
+            dynamic_axis = {0: "batch", 1: "sequence"}
+        return OrderedDict(
+            [
+                ("input_ids", dynamic_axis),
+                ("attention_mask", dynamic_axis),
+                ("token_type_ids", dynamic_axis),
+            ]
+        )

ixformer_sdk/contrib/transformers/models/t5/configuration_t5.py

@@ -0,0 +1,174 @@
+# coding=utf-8
+# Copyright 2020, The T5 Authors and HuggingFace Inc.
+#
+# 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.
+""" T5 model configuration"""
+from typing import Mapping
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.onnx import OnnxSeq2SeqConfigWithPast
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+T5_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "t5-small": "https://huggingface.co/t5-small/resolve/main/config.json",
+    "t5-base": "https://huggingface.co/t5-base/resolve/main/config.json",
+    "t5-large": "https://huggingface.co/t5-large/resolve/main/config.json",
+    "t5-3b": "https://huggingface.co/t5-3b/resolve/main/config.json",
+    "t5-11b": "https://huggingface.co/t5-11b/resolve/main/config.json",
+}
+
+
+class T5Config(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`T5Model`] or a [`TFT5Model`]. It is used to
+    instantiate a T5 model according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the T5
+    [t5-small](https://huggingface.co/t5-small) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Arguments:
+        vocab_size (`int`, *optional*, defaults to 32128):
+            Vocabulary size of the T5 model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`T5Model`] or [`TFT5Model`].
+        d_model (`int`, *optional*, defaults to 512):
+            Size of the encoder layers and the pooler layer.
+        d_kv (`int`, *optional*, defaults to 64):
+            Size of the key, query, value projections per attention head. The `inner_dim` of the projection layer will
+            be defined as `num_heads * d_kv`.
+        d_ff (`int`, *optional*, defaults to 2048):
+            Size of the intermediate feed forward layer in each `T5Block`.
+        num_layers (`int`, *optional*, defaults to 6):
+            Number of hidden layers in the Transformer encoder.
+        num_decoder_layers (`int`, *optional*):
+            Number of hidden layers in the Transformer decoder. Will use the same value as `num_layers` if not set.
+        num_heads (`int`, *optional*, defaults to 8):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        relative_attention_num_buckets (`int`, *optional*, defaults to 32):
+            The number of buckets to use for each attention layer.
+        relative_attention_max_distance (`int`, *optional*, defaults to 128):
+            The maximum distance of the longer sequences for the bucket separation.
+        dropout_rate (`float`, *optional*, defaults to 0.1):
+            The ratio for all dropout layers.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-6):
+            The epsilon used by the layer normalization layers.
+        initializer_factor (`float`, *optional*, defaults to 1):
+            A factor for initializing all weight matrices (should be kept to 1, used internally for initialization
+            testing).
+        feed_forward_proj (`string`, *optional*, defaults to `"relu"`):
+            Type of feed forward layer to be used. Should be one of `"relu"` or `"gated-gelu"`. T5v1.1 uses the
+            `"gated-gelu"` feed forward projection. Original T5 uses `"relu"`.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models).
+    """
+    model_type = "t5"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {
+        "hidden_size": "d_model",
+        "num_attention_heads": "num_heads",
+        "num_hidden_layers": "num_layers",
+    }
+
+    def __init__(
+        self,
+        vocab_size=32128,
+        d_model=512,
+        d_kv=64,
+        d_ff=2048,
+        num_layers=6,
+        num_decoder_layers=None,
+        num_heads=8,
+        relative_attention_num_buckets=32,
+        relative_attention_max_distance=128,
+        dropout_rate=0.1,
+        layer_norm_epsilon=1e-6,
+        initializer_factor=1.0,
+        feed_forward_proj="relu",
+        is_encoder_decoder=True,
+        use_cache=True,
+        pad_token_id=0,
+        eos_token_id=1,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.d_kv = d_kv
+        self.d_ff = d_ff
+        self.num_layers = num_layers
+        self.num_decoder_layers = (
+            num_decoder_layers if num_decoder_layers is not None else self.num_layers
+        )  # default = symmetry
+        self.num_heads = num_heads
+        self.relative_attention_num_buckets = relative_attention_num_buckets
+        self.relative_attention_max_distance = relative_attention_max_distance
+        self.dropout_rate = dropout_rate
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_factor = initializer_factor
+        self.feed_forward_proj = feed_forward_proj
+        self.use_cache = use_cache
+
+        act_info = self.feed_forward_proj.split("-")
+        self.dense_act_fn = act_info[-1]
+        self.is_gated_act = act_info[0] == "gated"
+
+        if len(act_info) > 1 and act_info[0] != "gated" or len(act_info) > 2:
+            raise ValueError(
+                f"`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."
+                "Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. "
+                "'gated-gelu' or 'relu'"
+            )
+
+        # for backwards compatibility
+        if feed_forward_proj == "gated-gelu":
+            self.dense_act_fn = "gelu_new"
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            is_encoder_decoder=is_encoder_decoder,
+            **kwargs,
+        )
+
+
+class T5OnnxConfig(OnnxSeq2SeqConfigWithPast):
+    @property
+    def inputs(self) -> Mapping[str, Mapping[int, str]]:
+        common_inputs = {
+            "input_ids": {0: "batch", 1: "encoder_sequence"},
+            "attention_mask": {0: "batch", 1: "encoder_sequence"},
+        }
+        if self.use_past:
+            common_inputs["attention_mask"][1] = "past_encoder_sequence + sequence"
+            common_inputs["decoder_input_ids"] = {0: "batch"}
+            common_inputs["decoder_attention_mask"] = {
+                0: "batch",
+                1: "past_decoder_sequence + sequence",
+            }
+        else:
+            common_inputs["decoder_input_ids"] = {0: "batch", 1: "decoder_sequence"}
+            common_inputs["decoder_attention_mask"] = {
+                0: "batch",
+                1: "decoder_sequence",
+            }
+
+        if self.use_past:
+            self.fill_with_past_key_values_(common_inputs, direction="inputs")
+
+        return common_inputs
+
+    @property
+    def default_onnx_opset(self) -> int:
+        return 13

ixformer_sdk/contrib/transformers/models/t5/ixformer_modeling_t5.py

@@ -0,0 +1,315 @@
+import torch
+from transformers.activations import NewGELUActivation
+
+import ixformer
+
+
+def self_attention_forward(
+    self,
+    hidden_states,
+    attention_mask=None,
+    position_bias=None,
+    layer_head_mask=None,
+    past_key_value=None,
+    use_cache=False,
+    output_attentions=False,
+):
+    assert output_attentions is False
+    assert layer_head_mask is None
+
+    normed_hidden_states = self.layer_norm(hidden_states)
+
+    if not hasattr(self, "qkv_weight"):
+        self.qkv_weight = torch.cat(
+            [
+                self.SelfAttention.q.weight,
+                self.SelfAttention.k.weight,
+                self.SelfAttention.v.weight,
+            ],
+            dim=0,
+        )
+        self.qkv_bias = None
+
+        del self.SelfAttention.q.weight
+        del self.SelfAttention.k.weight
+        del self.SelfAttention.v.weight
+
+    batch_size, seq_length = hidden_states.shape[:2]
+    real_seq_length = seq_length
+    if past_key_value is not None:
+        if len(past_key_value) != 2:
+            raise ValueError(
+                f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states"
+            )
+        real_seq_length += past_key_value[0].shape[2]
+    key_length = real_seq_length
+
+    def unshape(states):
+        """reshape"""
+        return (
+            states.transpose(1, 2)
+            .contiguous()
+            .view(batch_size, -1, self.SelfAttention.inner_dim)
+        )
+
+    qkv = ixformer.functions.linear(
+        normed_hidden_states, self.qkv_weight, self.qkv_bias
+    )
+
+    if past_key_value is not None:
+        pask_key, past_value = past_key_value
+        (
+            query_states,
+            key_states,
+            value_states,
+        ) = ixformer.functions.t5_split_qkv_update_kv_cache(
+            qkv,
+            pask_key,
+            past_value,
+            self.SelfAttention.n_heads,
+            self.SelfAttention.key_value_proj_dim,
+        )
+    else:
+        query_states, key_states, value_states = ixformer.functions.t5_split_qkv(
+            qkv, self.SelfAttention.n_heads, self.SelfAttention.key_value_proj_dim
+        )
+
+    if position_bias is None:
+        if not self.SelfAttention.has_relative_attention_bias:
+            position_bias = torch.zeros(
+                (1, self.SelfAttention.n_heads, real_seq_length, key_length),
+                device=query_states.device,
+                dtype=query_states.dtype,
+            )
+        else:
+            position_bias = self.SelfAttention.compute_bias(
+                real_seq_length, key_length, device=query_states.device
+            )
+
+        # if key and values are already calculated
+        # we want only the last query position bias
+        if past_key_value is not None:
+            position_bias = position_bias[:, :, -hidden_states.size(1) :, :]
+
+        if attention_mask is not None:
+            # (batch_size, n_heads, seq_length, key_length)
+            position_bias = position_bias + attention_mask
+
+    if self.SelfAttention.pruned_heads:
+        mask = torch.ones(position_bias.shape[1])
+        mask[list(self.pruned_heads)] = 0
+        position_bias_masked = position_bias[:, mask.bool()]
+    else:
+        position_bias_masked = position_bias
+
+    attn_output = ixformer.functions.ixinfer_flash_attn_pad(
+        query_states.contiguous(),
+        key_states.contiguous(),
+        value_states.contiguous(),
+        mask=position_bias_masked.float().contiguous(),
+        atten_scale=1,
+    )
+    attn_output = unshape(attn_output)
+
+    attn_output = self.SelfAttention.o(attn_output)
+
+    present_key_value_state = (
+        (key_states, value_states)
+        if (self.SelfAttention.is_decoder and use_cache)
+        else None
+    )
+    outputs = (attn_output,) + (present_key_value_state,) + (position_bias,)
+
+    if output_attentions:
+        outputs = outputs + (None,)
+    hidden_states = attn_output + hidden_states
+    outputs = (hidden_states,) + outputs[1:]
+
+    return outputs
+
+
+def cross_attention_forward(
+    self,
+    hidden_states,
+    key_value_states,
+    attention_mask=None,
+    position_bias=None,
+    layer_head_mask=None,
+    past_key_value=None,
+    use_cache=False,
+    query_length=None,
+    output_attentions=False,
+):
+
+    assert output_attentions is False
+    assert layer_head_mask is None
+
+    def unshape(states):
+        """reshape"""
+        return (
+            states.transpose(1, 2)
+            .contiguous()
+            .view(batch_size, -1, self.EncDecAttention.inner_dim)
+        )
+
+    normed_hidden_states = self.layer_norm(hidden_states)
+
+    # cross attn need key_value_states
+    assert key_value_states is not None
+    batch_size, seq_length = hidden_states.shape[:2]
+    real_seq_length = seq_length
+
+    if past_key_value is not None:
+        if len(past_key_value) != 2:
+            raise ValueError(
+                f"past_key_value should have 2 past states: keys and values. Got { len(past_key_value)} past states"
+            )
+        real_seq_length += (
+            past_key_value[0].shape[2] if query_length is None else query_length
+        )
+
+    key_length = (
+        real_seq_length if key_value_states is None else key_value_states.shape[1]
+    )
+    head_num, head_dim = (
+        self.EncDecAttention.n_heads,
+        self.EncDecAttention.key_value_proj_dim,
+    )
+
+    query_states = (
+        self.EncDecAttention.q(normed_hidden_states)
+        .view(batch_size, seq_length, head_num, head_dim)
+        .transpose(1, 2)
+        .contiguous()
+    )
+
+    if past_key_value is not None:
+        if past_key_value[0].shape[2] != key_value_states.shape[1]:
+            # checking that the `sequence_length` of the `past_key_value` is the same as
+            # the provided `key_value_states` to support prefix tuning
+            # cross-attn
+            # (batch_size, n_heads, seq_length, dim_per_head)
+            key_states = (
+                self.EncDecAttention.k(key_value_states)
+                .view(batch_size, key_length, head_num, head_dim)
+                .transpose(1, 2)
+            )
+            value_states = (
+                self.EncDecAttention.v(key_value_states)
+                .view(batch_size, key_length, head_num, head_dim)
+                .transpose(1, 2)
+            )
+        else:
+            # cross-attn
+            key_states = past_key_value[0]
+            value_states = past_key_value[1]
+    else:
+        key_states = (
+            self.EncDecAttention.k(key_value_states)
+            .view(batch_size, key_length, head_num, head_dim)
+            .transpose(1, 2)
+        )
+        value_states = (
+            self.EncDecAttention.v(key_value_states)
+            .view(batch_size, key_length, head_num, head_dim)
+            .transpose(1, 2)
+        )
+
+    if not query_states.is_contiguous():
+        query_states = query_states.contiguous()
+
+    # TODO: fix this bug
+    if not value_states.is_contiguous():
+        new_value_states = query_states.new_empty(value_states.shape)
+        new_value_states.copy_(value_states)
+        value_states = new_value_states
+    if not key_states.is_contiguous():
+        numel = torch.numel(key_states)
+        new_key_states = query_states.new_empty([numel * 2])[:numel].view(
+            *list(key_states.shape)
+        )
+        new_key_states.copy_(key_states)
+        key_states = new_key_states
+
+    if position_bias is None:
+        if not self.EncDecAttention.has_relative_attention_bias:
+            position_bias = torch.zeros(
+                (1, self.EncDecAttention.n_heads, real_seq_length, key_length),
+                device=query_states.device,
+                dtype=query_states.dtype,
+            )
+        else:
+            position_bias = self.EncDecAttention.compute_bias(
+                real_seq_length, key_length, device=query_states.device
+            )
+
+        # if key and values are already calculated
+        # we want only the last query position bias
+        if past_key_value is not None:
+            position_bias = position_bias[:, :, -hidden_states.size(1) :, :]
+
+        if attention_mask is not None:
+            # (batch_size, n_heads, seq_length, key_length)
+            position_bias = position_bias + attention_mask
+
+    if self.EncDecAttention.pruned_heads:
+        mask = torch.ones(position_bias.shape[1])
+        mask[list(self.pruned_heads)] = 0
+        position_bias_masked = position_bias[:, mask.bool()]
+    else:
+        position_bias_masked = position_bias
+
+    attn_output = ixformer.functions.ixinfer_flash_attn_pad(
+        query_states,
+        key_states.contiguous(),
+        value_states.contiguous(),
+        mask=position_bias_masked.float().contiguous(),
+        atten_scale=1,
+    )
+    attn_output = unshape(attn_output)
+
+    attn_output = self.EncDecAttention.o(attn_output)
+
+    present_key_value_state = (
+        (key_states, value_states)
+        if (self.EncDecAttention.is_decoder and use_cache)
+        else None
+    )
+    outputs = (attn_output,) + (present_key_value_state,) + (position_bias,)
+
+    if output_attentions:
+        outputs = outputs + (None,)
+    hidden_states = attn_output + hidden_states
+    outputs = (hidden_states,) + outputs[1:]
+
+    return outputs
+
+
+def dense_gated_act_dense_forward(self, hidden_states):
+    if isinstance(self.act, NewGELUActivation):
+        if not hasattr(self, "wi"):
+            self.wi = torch.cat([self.wi_1.weight, self.wi_0.weight], dim=0)
+            del self.wi_1
+            del self.wi_0
+        hidden_states = ixformer.functions.linear(hidden_states, self.wi, None)
+        hidden_states = ixformer.functions.gelu_and_mul(hidden_states)
+        hidden_states = ixformer.functions.linear(hidden_states, self.wo.weight, None)
+    else:
+        hidden_gelu = self.act(self.wi_0(hidden_states))
+        hidden_linear = self.wi_1(hidden_states)
+        hidden_states = hidden_gelu * hidden_linear
+
+        hidden_states = self.dropout(hidden_states)
+
+        # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32.
+        # See https://github.com/huggingface/transformers/issues/20287
+        # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None``
+        if (
+            isinstance(self.wo.weight, torch.Tensor)
+            and hidden_states.dtype != self.wo.weight.dtype
+            and self.wo.weight.dtype != torch.int8
+        ):
+            hidden_states = hidden_states.to(self.wo.weight.dtype)
+
+        hidden_states = self.wo(hidden_states)
+    return hidden_states

ixformer_sdk/contrib/vllm/layers/__init__.py

@@ -0,0 +1,30 @@
+from .llama import forward_smoothquant
+from .mixtral import mixtral_decoder_layer_forward
+
+SUPPORT_REPLACE_METHOD = {
+    "llama": forward_smoothquant,
+}
+
+SUPPORT_REPLACE_LAYER = {
+    "llama": None,
+}
+
+
+def get_replace_forward(name: str):
+    try:
+        method = SUPPORT_REPLACE_METHOD[name]
+    except:
+        raise ValueError(
+            f"Only support replace names: {SUPPORT_REPLACE_METHOD.keys()}, but got {name}"
+        )
+    return method
+
+
+def get_replace_layer(name: str):
+    try:
+        layer = SUPPORT_REPLACE_LAYER[name]
+    except:
+        raise ValueError(
+            f"Only support replace names: {SUPPORT_REPLACE_LAYER.keys()}, but got {name}"
+        )
+    return layer

ixformer_sdk/contrib/vllm/layers/llama.py

@@ -0,0 +1,100 @@
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import torch
+from transformers import LlamaConfig
+
+from vllm.attention import AttentionMetadata
+from vllm.config import CacheConfig
+from vllm.distributed import tensor_model_parallel_all_reduce
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.models.llama import LlamaDecoderLayer as VllmLlamaDecoderLayer
+
+import vllm._custom_ops as ops
+# from ..overlap_comm import DecoderLayerOverlapComm, get_overlap_linear_method
+
+
+# This method is needed for support smoothquant no overlap forward
+def forward_smoothquant(
+    input_ids: Optional[torch.Tensor],
+    positions: torch.Tensor,
+    kv_caches: List[torch.Tensor],
+    attn_metadata: AttentionMetadata,
+    inputs_embeds: Optional[torch.Tensor] = None,
+    self = None, # will be set by partial
+) -> torch.Tensor:
+    dtype = self.dtype
+    
+    def forward_smoothquant_mlp(self,x,scales):
+        # gate_up_proj
+        # Int8 Matrix multiply.
+        bias = self.gate_up_proj.bias if not self.gate_up_proj.skip_bias_add else None
+        gate_up = ops.w8a8(x, self.gate_up_proj.weight, scales, self.gate_up_proj.weight_scales, dtype)
+        if bias:
+            gate_up += bias
+        
+        # act_fun
+        x, scales = ops.silu_and_mul_smoothquant(gate_up, self.down_proj.smooth_scales)
+        
+        # down_proj
+        output_parallel = ops.w8a8(x, self.down_proj.weight, scales, self.down_proj.weight_scales, dtype)
+        if self.down_proj.reduce_results and self.down_proj.tp_size > 1:
+            output = tensor_model_parallel_all_reduce(output_parallel)
+        else:
+            output = output_parallel
+
+        if not self.down_proj.skip_bias_add:
+            output = output + self.down_proj.bias if self.down_proj.bias is not None else output
+
+        return output
+    
+    def forward_smoothquant_attn(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+        scales: torch.Tensor,
+    ) -> torch.Tensor:
+        # qkv proj
+        bias = self.qkv_proj.bias if not self.qkv_proj.skip_bias_add else None
+        
+        qkv = ops.w8a8(hidden_states, self.qkv_proj.weight, scales, self.qkv_proj.weight_scales, dtype)
+        if bias:
+            qkv += bias
+        
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k = self.rotary_emb(positions, q, k)
+        attn_output = self.attn(q, k, v, kv_cache, attn_metadata)
+        output, _ = self.o_proj(attn_output)  # TODO
+        return output
+    
+    if inputs_embeds is not None:
+        hidden_states = inputs_embeds
+    else:
+        hidden_states = self.get_input_embeddings(input_ids)
+    residual = None
+    for i in range(len(self.layers)):
+        layer = self.layers[i]
+        if residual is None:
+            residual = hidden_states
+            hidden_states, scales = ops.rms_norm_smoothquant(hidden_states,layer.input_layernorm.weight,layer.input_layernorm.variance_epsilon, layer.self_attn.qkv_proj.smooth_scales)
+        else:
+            hidden_states, residual, scales = ops.fused_add_rms_norm_smoothquant(hidden_states, residual, layer.input_layernorm.weight, layer.input_layernorm.variance_epsilon, layer.self_attn.qkv_proj.smooth_scales)
+        
+        hidden_states = forward_smoothquant_attn(
+            layer.self_attn,
+            positions=positions,
+            hidden_states=hidden_states,
+            kv_cache=kv_caches[i],
+            attn_metadata=attn_metadata,
+            scales=scales,
+            )
+
+        # Fully Connected
+        hidden_states, residual, scales = ops.fused_add_rms_norm_smoothquant(hidden_states, residual, layer.post_attention_layernorm.weight, layer.post_attention_layernorm.variance_epsilon, layer.mlp.gate_up_proj.smooth_scales)
+        
+        hidden_states = forward_smoothquant_mlp(layer.mlp, hidden_states, scales)
+
+    hidden_states, _ = self.norm(hidden_states, residual)
+    return hidden_states
+

ixformer_sdk/contrib/vllm/layers/mixtral.py

@@ -0,0 +1,331 @@
+import functools
+from typing import Dict, Optional, Tuple
+
+import ixformer.inference.functions as ixf
+import torch
+
+
+def mixtral_decoder_layer_forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata,
+    residual: Optional[torch.Tensor],
+) -> torch.Tensor:
+    if self.use_int_w8a8:
+        return w8a8_forward(
+            self, positions, hidden_states, kv_cache, attn_metadata, residual
+        )
+    else:
+        return original_forward(
+            self, positions, hidden_states, kv_cache, attn_metadata, residual
+        )
+
+
+def original_forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata,
+    residual: Optional[torch.Tensor],
+) -> torch.Tensor:
+    # Self Attention
+    if residual is None:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+    else:
+        hidden_states, residual = self.input_layernorm(hidden_states, residual)
+    hidden_states = self.self_attn(
+        positions=positions,
+        hidden_states=hidden_states,
+        kv_cache=kv_cache,
+        attn_metadata=attn_metadata,
+    )
+
+    # Fully Connected
+    hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+    hidden_states = self.block_sparse_moe(hidden_states)
+    return hidden_states, residual
+
+
+def dynamic_scaled_int8_quant(x):
+    m, k = x.shape
+    i8_x = x.new_empty([m, k], dtype=torch.int8, device="cuda")
+    i8_scales = torch.empty([m], dtype=torch.float32, device="cuda")
+    ixf.dynamic_scaled_int8_quant(i8_x, x, i8_scales)
+    return i8_x, i8_scales
+
+
+def dynamic_w8a8(x, i8_weight, weight_scale):
+    i8_x, i8_scale = dynamic_scaled_int8_quant(x)
+    m, k = x.shape
+    k, n = i8_weight.shape
+    output = x.new_empty([m, n], dtype=x.dtype, device="cuda")
+    ixf.w8a8(
+        i8_x,
+        i8_weight.transpose(0, 1),
+        i8_scale,
+        weight_scale,
+        output=output,
+        out_dtype=x.dtype,
+    )
+    return output
+
+
+def fused_rms_norm_quant_linear(
+    self,
+    hidden_states,
+    ln_weight,
+    eps,
+    linear_weight,
+    linear_weight_scale,
+    residual=None,
+):
+    # lower rouge
+    # if residual is None:
+    #     residual = hidden_states
+    #     i8_hidden_states, _, i8_scales = ixf.residual_rms_norm_dynamic_int8(
+    #         input=hidden_states,
+    #         weight=ln_weight,
+    #         residual=None,
+    #         eps=eps,
+    #     )
+    # else:
+    #     i8_hidden_states, residual, i8_scales = ixf.residual_rms_norm_dynamic_int8(
+    #         input=hidden_states,
+    #         weight=ln_weight,
+    #         residual=residual,
+    #         eps=eps,
+    #     )
+
+    if residual is None:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+    else:
+        hidden_states, residual = self.input_layernorm(hidden_states, residual)
+    i8_hidden_states, i8_scales = dynamic_scaled_int8_quant(hidden_states)
+
+    qkv = hidden_states.new_empty(hidden_states.shape[0], linear_weight.shape[1])
+    ixf.w8a8(
+        i8_hidden_states,
+        linear_weight.transpose(0, 1),
+        i8_scales,
+        linear_weight_scale,
+        output=qkv,
+        out_dtype=hidden_states.dtype,
+    )
+    return qkv, residual
+
+
+def attention(qkv, positions, kv_cache, attn_metadata, self_attn):
+    q, k, v = qkv.split(
+        [self_attn.q_size, self_attn.kv_size, self_attn.kv_size], dim=-1
+    )
+    q, k = self_attn.rotary_emb(positions, q, k)
+    attn_output = self_attn.attn(q, k, v, kv_cache, attn_metadata)
+    return attn_output
+
+
+def fused_rms_norm_attention(
+    self,
+    hidden_states,
+    ln_weight,
+    eps,
+    positions,
+    kv_cache,
+    attn_metadata,
+    self_attn,
+    residual=None,
+):
+    hidden_states, residual = fused_rms_norm_quant_linear(
+        self,
+        hidden_states,
+        ln_weight,
+        eps,
+        self_attn.qkv_proj.weight,
+        self_attn.qkv_proj.weight_scale,
+        residual,
+    )
+    hidden_states = attention(
+        hidden_states, positions, kv_cache, attn_metadata, self_attn
+    )
+
+    hidden_states = dynamic_w8a8(
+        hidden_states, self_attn.o_proj.weight, self_attn.o_proj.weight_scale
+    )
+    # hidden_states,_ = self_attn.o_proj(hidden_states) # quant+linear+allreduce
+    return hidden_states, residual
+
+
+def w8a8_forward(
+    self,
+    positions: torch.Tensor,
+    hidden_states: torch.Tensor,
+    kv_cache: torch.Tensor,
+    attn_metadata,
+    residual: Optional[torch.Tensor],
+) -> torch.Tensor:
+
+    # qkv,_ = self.self_attn.qkv_proj(hidden_states)
+    hidden_states, residual = fused_rms_norm_attention(
+        self,
+        hidden_states,
+        self.input_layernorm.weight,
+        self.input_layernorm.variance_epsilon,
+        positions,
+        kv_cache,
+        attn_metadata,
+        self.self_attn,
+        residual,
+    )
+
+    # allreduce
+    tp_size = self.block_sparse_moe.experts.tp_size
+    if tp_size > 1:
+        from vllm.distributed import tensor_model_parallel_all_reduce
+
+        hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+
+    # rms norm
+    hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+    # moe
+    hidden_states = fused_moe(
+        hidden_states,
+        self.block_sparse_moe.gate.weight,
+        top_k=self.block_sparse_moe.experts.top_k,
+        w1=self.block_sparse_moe.experts.w13_weight,
+        w2=self.block_sparse_moe.experts.w2_weight,
+        w1_scale=self.block_sparse_moe.experts.w13_weight_scale,
+        w2_scale=self.block_sparse_moe.experts.w2_weight_scale,
+    )
+
+    # allreduce
+    if tp_size > 1:
+        from vllm.distributed import tensor_model_parallel_all_reduce
+
+        hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+
+    return hidden_states, residual
+
+
+def fused_experts(hidden_states, router_logits, top_k, w1, w2, w1_scale, w2_scale):
+
+    """
+    Args:
+        hidden_states:          (num_tokens, k)                 dtype
+        router_logits:          (num_tokens, num_experts)       torch.float32
+        top_k                                                   int
+        w1:                     (num_experts, 2n, k)            torch.int8
+        w2:                     (num_experts, k, n)             torch.int8
+        w1_scale:               (num_experts, 2n)               torch.float32
+        w2_scale:               (num_experts, k)                torch.float32
+    Returns
+        final_hidden_states:    (num_tokens, k)                 dtype
+    """
+
+    # topk_weight: (num_tokens, top_k) torch.float32
+    # topk_ids:    (num_tokens, top_k) torch.int32
+    topk_weight, topk_ids = ixf.moe_topk_softmax(
+        gating_output=router_logits,
+        topk=top_k,
+        renormalize=True,
+    )
+
+    dtype = hidden_states.dtype
+    num_tokens, num_experts = router_logits.shape
+    expand_tokens = num_tokens * top_k
+
+    (
+        src_to_dst,
+        sorted_token_ids,
+        expert_sizes_gpu,
+        expert_sizes_cpu,
+    ) = ixf.moe_compute_token_index(
+        topk_ids=topk_ids,
+        num_experts=num_experts,
+    )
+    expert_sizes_cpu = expert_sizes_gpu.cpu()
+
+    # expand + reorder + quant
+    # i8_hidden_states: (expand_tokens, k) torch.int8
+    i8_hidden_states, a_scale = ixf.moe_expand_input_dynamic_scaled_int8(
+        hidden_states=hidden_states,
+        dst_to_src=sorted_token_ids,
+        dst_tokens=expand_tokens,
+        topk=top_k,
+        src_to_dst=src_to_dst,
+        topk_ids=None,  # use smooth quant
+        smooth_scales=None,  # use smooth quant
+    )
+
+    # w8a8 group gemm 1
+    # pt_output_1: (expand_tokens, 2n) dtype
+    pt_output_1 = ixf.moe_w8a8_group_gemm(
+        input=i8_hidden_states,
+        weight=w1,
+        i_scales=a_scale,
+        w_scales=w1_scale,
+        output_dtype=dtype,
+        tokens_per_experts=expert_sizes_cpu,
+        dst_to_src=None,
+        format="TN",
+    )
+
+    # act + quant
+    # pt_output_2: (expand_tokens, n) torch.int8
+    pt_output_2, a2_scale = ixf.activation_dynamic_scaled_int8(
+        input=pt_output_1,
+        bias=None,  # add gemm bias
+        smooth_scales=None,  # use smooth quant
+        dst_to_src=sorted_token_ids,
+        topk_ids=None,  # add gemm bias or use smooth quant
+        act_type="swiglu",
+    )
+
+    # w8a8 group gemm 2 + reorder
+    # pt_output_3: (expand_tokens, k) dtype
+    pt_output_3 = ixf.moe_w8a8_group_gemm(
+        input=pt_output_2,
+        weight=w2,
+        i_scales=a2_scale,
+        w_scales=w2_scale,
+        output_dtype=dtype,
+        tokens_per_experts=expert_sizes_cpu,
+        dst_to_src=sorted_token_ids,
+        format="TN",
+    )
+
+    # mul + reduce_sum
+    # final_hidden_states: (num_tokens, k)
+    final_hidden_states = ixf.moe_output_reduce_sum(
+        input=pt_output_3.view(num_tokens, top_k, -1),
+        topk_weight=topk_weight,
+    )
+
+    return final_hidden_states
+
+
+def fused_moe(hidden_states, gate_weight, top_k, w1, w2, w1_scale, w2_scale):
+    orig_shape = hidden_states.shape
+    hidden_size = hidden_states.shape[-1]
+
+    hidden_states = hidden_states.view(-1, hidden_size)
+
+    # router_logits: (num_tokens, n_experts)
+    # gate_weight: fp16
+    router_logits = ixf.linear(hidden_states, gate_weight)
+    router_logits = router_logits.to(torch.float32)
+
+    final_hidden_states = fused_experts(
+        hidden_states,
+        router_logits,
+        top_k,
+        w1,
+        w2,
+        w1_scale,
+        w2_scale,
+    )
+
+    return final_hidden_states.view(orig_shape)

ixformer_sdk/contrib/vllm/quantize/__init__.py

@@ -0,0 +1,14 @@
+from .smoothquant import smoothquant_prepare_quantize,smoothquant_export_quantized_weights
+from .w8a16 import w8a16_prepare_quantize,w8a16_export_quantized_weights
+
+SUPPORT_METHOD = {
+    "smoothquant": [smoothquant_prepare_quantize,smoothquant_export_quantized_weights],
+    "w8a16": [w8a16_prepare_quantize,w8a16_export_quantized_weights],
+}
+
+def get_quantize_method(method_name:str):
+    try:
+        method = SUPPORT_METHOD[method_name]
+    except:
+        raise ValueError(f"Only support quantization methods: {SUPPORT_METHOD.keys()}, but got {method_name}")
+    return method

ixformer_sdk/contrib/vllm/quantize/smoothquant.py

@@ -0,0 +1,407 @@
+import os
+
+import torch
+
+
+def smoothquant_prepare_quantize(self, quant_params={}):
+    model = self.model_runner.model
+
+    def update_act_scales(act_scales, x):
+        # 动态统计每次输入的最大值
+        hidden_dim = x.shape[-1]
+        x = x.view(-1, hidden_dim).abs().detach()
+        # [k]
+        comming_max = torch.max(x, dim=0, keepdim=True)[0].float()
+
+        if act_scales is None:
+            act_scales = comming_max
+        else:
+            act_scales = torch.max(act_scales, comming_max)
+        return act_scales
+
+    from functools import partial
+
+    from vllm.model_executor.layers.linear import (
+        ColumnParallelLinear,
+        MergedColumnParallelLinear,
+        QKVParallelLinear,
+        RowParallelLinear,
+    )
+
+    def new_forward(input_, m, raw_forward):
+        if not hasattr(m, "act_scales"):
+            m.act_scales = None
+        m.act_scales = update_act_scales(m.act_scales, input_)
+        return raw_forward(input_)
+
+    for name, m in model.named_modules():
+        if (
+            isinstance(m, QKVParallelLinear)
+            or isinstance(m, RowParallelLinear)
+            or isinstance(m, MergedColumnParallelLinear)
+            or isinstance(m, ColumnParallelLinear)
+        ):
+            m.forward = partial(new_forward, m=m, raw_forward=m.forward)
+
+def smoothquant_export_quantized_weights(self, save_path, quant_params={}):
+    gb_per_file = quant_params.get("filesize_limit", None)
+    smooth_alpha = quant_params.get("smooth_alpha", 0.5)
+    dynamic_quant_type = quant_params.get("dynamic_quant_type", "gpu")
+    assert dynamic_quant_type in ["gpu","cpu","kernel"]
+    if self.rank == 0:
+        print(f"set smooth_alpha={smooth_alpha}")
+        print(f"use quantize weight type: {dynamic_quant_type}")
+    
+    import ixformer._C as ops
+    def per_token_quant_8bit(weight):
+        # weight: [m,k]
+        dtype = weight.dtype
+        i8_weight = weight
+        scale = i8_weight.abs().max(dim=-1, keepdim=True)[0] / 127
+        i8_weight = i8_weight / scale.to(dtype)
+        i8_weight = torch.clamp(torch.round(i8_weight), -128, 127).to(torch.int8)
+        return i8_weight, scale.float()
+
+    def smooth_quant_weight_gpu_cpu(weight, act_scale, alpha=0.5, device="cpu"):
+        device = torch.device("cpu") if device == "cpu" else weight.device
+        ori_dtype = weight.dtype
+        # [1, k]
+        act_scale = act_scale.float().to(device).view(1, -1)
+        weight = weight.to(device)
+        # [1, k]
+        weight_scale = weight.abs().max(dim=0, keepdim=True)[0].float()
+        if alpha == -1:
+            smooth_scales = torch.ones_like(act_scale)
+        else:
+            smooth_scales = act_scale.pow(alpha) / weight_scale.pow(1 - alpha).clamp(
+                min=1e-5
+            )
+        weight = weight * smooth_scales.to(ori_dtype)
+        i8_weight, weight_scales = per_token_quant_8bit(weight)
+        # 为了可以使用 input * smooth_scales
+        if alpha == -1:
+            smooth_scales = torch.ones_like(act_scale)
+        else:
+            smooth_scales = weight_scale.pow(1 - alpha) / act_scale.pow(alpha).clamp(
+                min=1e-5
+            )
+        return i8_weight, weight_scales, smooth_scales.to(ori_dtype)
+    
+    def smooth_quant_weight_kernel(weight, act_scale, alpha=0.5):
+        output = torch.zeros_like(weight,dtype=torch.int8)
+        weight_scales = torch.zeros(weight.shape[:-1],dtype=torch.float, device=weight.device)
+        weight_max = torch.zeros(weight.shape[-1],dtype=torch.float, device=weight.device)
+        smooth_scales = torch.zeros(weight.shape[-1],dtype=weight.dtype, device=weight.device)
+        ops.infer.weight_quant_smoothquant(
+            weight, act_scale, alpha, output, weight_scales, smooth_scales, weight_max
+        )
+        return output, weight_scales.view(-1,1), smooth_scales.view(1,-1)
+    
+    def smooth_quant_weight(weight, act_scale, alpha=0.5):
+        if dynamic_quant_type == "kernel":
+            return smooth_quant_weight_kernel(weight,act_scale,alpha)
+        else:
+            return smooth_quant_weight_gpu_cpu(weight,act_scale,alpha,dynamic_quant_type)
+
+    model = self.model_runner.model
+
+    from vllm.distributed import (
+        tensor_model_parallel_all_gather,
+        tensor_model_parallel_all_reduce,
+        get_tensor_model_parallel_world_size
+    )
+    from vllm.model_executor.layers.linear import (
+        ColumnParallelLinear,
+        MergedColumnParallelLinear,
+        QKVParallelLinear,
+        RowParallelLinear,
+    )
+    from vllm.model_executor.layers.vocab_parallel_embedding import (
+        ParallelLMHead,
+        VocabParallelEmbedding,
+    )
+    from vllm.model_executor.models.falcon import FalconForCausalLM
+
+    for name, m in model.named_modules():
+        if isinstance(m, VocabParallelEmbedding):
+            # weight shape: [vocab_size // tp, embedding_dim]
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
+                print(f"merged: {name}, shape={m.weight.shape}")
+            
+        elif isinstance(m, ParallelLMHead):
+            # weight shape: [vocab_size // tp, embedding_dim]
+            # bias shape: [vocab_size // tp]
+            if m.bias is not None:
+                bias = tensor_model_parallel_all_gather(m.bias, dim=0)
+                bias = bias[:m.org_vocab_size].contiguous()
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(bias.cpu(), requires_grad=False)
+            
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
+                print(f"merged: {name}, shape={m.weight.shape}")
+        
+        elif isinstance(m, QKVParallelLinear):
+            # weight shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp, hidden_size]
+            # bias shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp]
+            if self.parallel_config.world_size > 1:
+                total_q_hidden_size = m.total_num_heads * m.head_size
+                partial_q_hidden_size = m.num_heads * m.head_size
+                total_kv_hidden_size = m.total_num_kv_heads * m.head_size
+                partial_kv_hidden_size = m.num_kv_heads * m.head_size
+                
+                if m.bias is not None:
+                    # TODO do not support padding..
+                    bias_tenosr = m.bias.new_zeros(total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size)
+                    q_bias = bias_tenosr[:total_q_hidden_size][self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
+                    k_bias = bias_tenosr[total_q_hidden_size:total_q_hidden_size+total_kv_hidden_size]\
+                                [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+                    v_bias = bias_tenosr[total_q_hidden_size+total_kv_hidden_size:]\
+                                [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+                    
+                    q_bias[:] = m.bias[:partial_q_hidden_size]
+                    k_bias[:] = m.bias[partial_q_hidden_size:partial_q_hidden_size+partial_kv_hidden_size]
+                    v_bias[:] = m.bias[partial_q_hidden_size+partial_kv_hidden_size:]
+                    
+                    bias_tensor = tensor_model_parallel_all_reduce(bias_tenosr)
+                    if self.is_driver_worker:
+                        m.bias = torch.nn.Parameter(bias_tensor.cpu(), requires_grad=False)
+
+                q_tensor = m.weight.new_zeros(m.total_num_heads * m.head_size, m.weight.shape[1])
+                k_tensor = m.weight.new_zeros(m.total_num_kv_heads * m.head_size, m.weight.shape[1])
+                v_tensor = m.weight.new_zeros(m.total_num_kv_heads * m.head_size, m.weight.shape[1])
+                
+                q_in_weight = m.weight[:-m.num_kv_heads * m.head_size * 2]
+                k_in_weight = m.weight[-m.num_kv_heads * m.head_size * 2:-m.num_kv_heads * m.head_size]
+                v_in_weight = m.weight[-m.num_kv_heads * m.head_size:]
+                
+                if getattr(m,"start_idx",None) is not None:
+                    start_idx = getattr(m,"start_idx")
+                    weight_end_idx = m.num_heads * m.head_size if not getattr(m,"is_padding") else (m.num_heads - 1) * m.head_size
+                    end_idx = start_idx + weight_end_idx
+                else:
+                    start_idx = self.rank * m.num_heads * m.head_size
+                    weight_end_idx = m.num_heads * m.head_size
+                    end_idx = start_idx + weight_end_idx
+                assert q_tensor[start_idx:end_idx,:].shape == q_in_weight[:weight_end_idx, :].shape
+                q_tensor[start_idx:end_idx,:] = q_in_weight[:weight_end_idx, :]
+                
+                if m.num_kv_head_replicas > 1:
+                    if self.rank % m.num_kv_head_replicas == 0:
+                        rank = self.rank // m.num_kv_head_replicas
+                        k_tensor[rank * m.num_kv_heads * m.head_size:(rank+1) * m.num_kv_heads * m.head_size] = k_in_weight
+                        v_tensor[rank * m.num_kv_heads * m.head_size:(rank+1) * m.num_kv_heads * m.head_size] = v_in_weight
+                else:
+                    k_tensor[self.rank * m.num_kv_heads * m.head_size:(self.rank+1) * m.num_kv_heads * m.head_size] = k_in_weight
+                    v_tensor[self.rank * m.num_kv_heads * m.head_size:(self.rank+1) * m.num_kv_heads * m.head_size] = v_in_weight
+                
+                q_tensor = tensor_model_parallel_all_reduce(q_tensor)
+                k_tensor = tensor_model_parallel_all_reduce(k_tensor)
+                v_tensor = tensor_model_parallel_all_reduce(v_tensor)
+                
+                if isinstance(model, FalconForCausalLM):
+                    num_query_heads_per_kv_head = (
+                        m.total_num_heads // m.total_num_kv_heads
+                    )
+                    q_tensor = q_tensor.view(
+                        m.total_num_kv_heads,
+                        num_query_heads_per_kv_head,
+                        m.head_size,
+                        -1,
+                    )
+                    k_tensor = k_tensor.view(m.total_num_kv_heads, 1, m.head_size, -1)
+                    v_tensor = v_tensor.view(m.total_num_kv_heads, 1, m.head_size, -1)
+                    weight_tensor = torch.cat(
+                        [q_tensor, k_tensor, v_tensor], dim=1
+                    ).view(-1, m.hidden_size)
+                else:
+                    weight_tensor = torch.cat([q_tensor, k_tensor, v_tensor])
+                assert (
+                    weight_tensor.shape[0]
+                    == total_q_hidden_size + total_kv_hidden_size * 2
+                )
+                assert weight_tensor.shape[1] == m.hidden_size
+
+            else:
+                weight_tensor = m.weight
+                if m.bias is not None and self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+            
+            if self.is_driver_worker:
+                i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
+                    weight_tensor, m.act_scales, smooth_alpha
+                )
+                m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
+                m.weight_scales = torch.nn.Parameter(
+                    weight_scales.cpu(), requires_grad=False
+                )
+                m.smooth_scales = torch.nn.Parameter(
+                    smooth_scales.cpu(), requires_grad=False
+                )
+                print(f"Quantized: {name}")
+            
+        elif isinstance(m, MergedColumnParallelLinear):
+            if self.parallel_config.world_size > 1:
+                # weight shape: [intermediate_size // tp * 2, hidden_size]
+                # bias shape: [intermediate_size // tp * 2]
+                output_sizes = m.output_sizes
+                output_size = sum(output_sizes)
+                partial_output_sizes = [
+                    i // self.parallel_config.world_size for i in output_sizes
+                ]
+
+                if m.bias is not None:
+                    index_start = 0
+                    partial_index_start = 0
+                    bias_tenosr = m.bias.new_zeros(output_size)
+                    for i in range(len(output_sizes)):
+                        index_out = index_start + output_sizes[i]
+                        sub_bias_tensor = bias_tenosr[index_start:index_out]
+                        partial_size = partial_output_sizes[i]
+                        sub_bias_tensor[self.rank * partial_size:(self.rank+1) * partial_size] = m.bias[partial_index_start:partial_index_start+partial_size]
+                        
+                        index_start += output_sizes[i]
+                        partial_index_start += partial_size
+                    bias_tenosr = tensor_model_parallel_all_reduce(bias_tenosr)
+                    if self.is_driver_worker:
+                        m.bias = torch.nn.Parameter(bias_tenosr, requires_grad=False)
+                
+                weight_tensor = m.weight.new_zeros(output_size, m.input_size)
+
+                idx_out_start = 0
+                idx_partial_satrt = 0
+                for i in range(len(output_sizes)):
+                    idx_out_end = idx_out_start + output_sizes[i]
+                    sub_weight_tensor = weight_tensor[idx_out_start:idx_out_end]
+                    partial_size = partial_output_sizes[i]
+                    sub_weight_tensor[
+                        self.rank * partial_size : (self.rank + 1) * partial_size
+                    ] = m.weight[idx_partial_satrt : idx_partial_satrt + partial_size]
+
+                    idx_out_start += output_sizes[i]
+                    idx_partial_satrt += partial_size
+                weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
+            else:
+                weight_tensor = m.weight
+                if m.bias is not None and self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+            
+            if self.is_driver_worker:
+                i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
+                    weight_tensor, m.act_scales, smooth_alpha
+                )
+                m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
+                m.weight_scales = torch.nn.Parameter(
+                    weight_scales.cpu(), requires_grad=False
+                )
+                m.smooth_scales = torch.nn.Parameter(
+                    smooth_scales.cpu(), requires_grad=False
+                )
+                print(f"Quantized: {name}")
+
+        elif isinstance(m, ColumnParallelLinear):
+            # weight shape: [some_dim // tp, hidden_size] // for this Linear, some_dim mostly is hidden_size * 4
+            # bias shape: [some_dim // tp]
+            if m.bias is not None:
+                bias_tenosr = tensor_model_parallel_all_gather(m.bias, dim=0)
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(bias_tenosr.cpu(), requires_grad=False)
+
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+
+            if self.is_driver_worker:
+                i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
+                   weight_tensor, m.act_scales, smooth_alpha
+                )
+                m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
+                m.weight_scales = torch.nn.Parameter(
+                    weight_scales.cpu(), requires_grad=False
+                )
+                m.smooth_scales = torch.nn.Parameter(
+                    smooth_scales.cpu(), requires_grad=False
+                )
+                print(f"Quantized: {name}")
+
+        elif isinstance(m, RowParallelLinear):
+            # weight shape: [hidden_size, some_dim // tp] // for this Linear, some_dim mostly is hidden_size * 4 or intermediate_size
+            # bias shape: [hidden_size]
+            if m.bias is not None:
+                bias_tensor = tensor_model_parallel_all_gather(m.bias, dim=-1)
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+            
+            if getattr(m,"start_idx", None) is not None:
+                start_idx = getattr(m,"start_idx")
+                end_idx = start_idx + (m.input_size_per_partition if not getattr(m,"is_padding") else (m.input_size_per_partition - m.padding_size))
+                weight_end_idx = m.input_size_per_partition if not getattr(m,"is_padding") else (m.input_size_per_partition - m.padding_size)
+            else:
+                start_idx = m.input_size_per_partition * self.rank
+                end_idx = start_idx + m.input_size_per_partition
+                weight_end_idx = m.input_size_per_partition
+            
+            act_scales = m.act_scales.new_zeros(m.input_size)
+            assert act_scales[start_idx:end_idx].shape == m.act_scales.view(-1)[:weight_end_idx].shape
+            act_scales[start_idx:end_idx] = m.act_scales.view(-1)[:weight_end_idx]
+            act_scales = tensor_model_parallel_all_reduce(act_scales)
+            m.act_scales = act_scales
+            
+            weight_tensor = m.weight.new_zeros(m.weight.shape[0],m.input_size)
+            assert  weight_tensor[:,start_idx:end_idx].shape == m.weight[:,:weight_end_idx].shape
+            weight_tensor[:,start_idx:end_idx] = m.weight[:,:weight_end_idx]
+            weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
+            
+            if self.is_driver_worker:
+                i8_weight, weight_scales, smooth_scales = smooth_quant_weight(
+                    weight_tensor, m.act_scales, smooth_alpha
+                )
+                smooth_scales = smooth_scales.view(1,-1)
+                m.weight = torch.nn.Parameter(i8_weight.cpu(), requires_grad=False)
+                m.weight_scales = torch.nn.Parameter(
+                    weight_scales.cpu(), requires_grad=False
+                )
+                m.smooth_scales = torch.nn.Parameter(
+                    smooth_scales.cpu(), requires_grad=False
+                )
+                print(f"Quantized: {name}")
+        else:
+            pass
+        
+        torch.cuda.empty_cache()
+
+    # save weights
+    if self.is_driver_worker:
+        from safetensors.torch import save_file
+        
+        tensors = {}
+        saved = False
+        count = 0
+        size_in_bytes = 0
+        
+        tensors = {}
+        for name, weight in model.named_parameters():
+            if "act_scales" in name:
+                continue
+            # skip lm_head_weight if needed..
+            if "lm_head" in name and model.config.tie_word_embeddings:
+                continue
+            tensors[name] = weight
+            
+            saved = False
+            if gb_per_file is not None and size_in_bytes >= gb_per_file * 1024 * 1024 * 1024:
+                weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))
+                save_file(tensors, weight_path)
+                print(f"The quantified weights were successfully saved in {weight_path}.")
+                tensors.clear()
+                saved = True
+                count += 1
+                size_in_bytes = 0
+        
+        if not saved: 
+            weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))   
+            save_file(tensors, weight_path)
+            print(f"The quantified weights were successfully saved in {weight_path}.")

ixformer_sdk/contrib/vllm/quantize/w8a16.py

@@ -0,0 +1,233 @@
+import os
+
+import torch
+
+def w8a16_prepare_quantize(self, quant_params={}):
+    # We need do nothing in here
+    pass
+
+
+def w8a16_export_quantized_weights(self, save_path, quant_params={}):
+    gb_per_file = quant_params.get("filesize_limit", None)
+    int8_min = -127
+
+    def w8a16_quantization(weight):
+        # all weights should be [output,input], otherwise, we may get an wrong weight and scale...
+        scale = torch.abs(weight).max(dim=-1)[0] / 127.0
+        int8_weight = torch.clamp(weight / scale.view(-1,1),min=int8_min,max=127).to(torch.int8).contiguous()
+        scale = scale.view(1,-1).contiguous()
+        return int8_weight, scale
+
+
+    model = self.model_runner.model
+
+    from vllm.distributed.communication_op import (
+        tensor_model_parallel_all_gather,
+        tensor_model_parallel_all_reduce,
+    )
+    from vllm.model_executor.layers.linear import (
+        ColumnParallelLinear,
+        MergedColumnParallelLinear,
+        QKVParallelLinear,
+        RowParallelLinear,
+    )
+    from vllm.model_executor.layers.vocab_parallel_embedding import (
+        ParallelLMHead,
+        VocabParallelEmbedding,
+    )
+
+    for name, m in model.named_modules():
+        if isinstance(m, VocabParallelEmbedding):
+            # weight shape: [vocab_size // tp, embedding_dim]
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
+                print(f"merged: {name}, shape={m.weight.shape}")
+        
+        elif isinstance(m, ParallelLMHead):
+            # weight shape: [vocab_size // tp, embedding_dim]
+            # bias shape: [vocab_size // tp]
+            if m.bias is not None:
+                bias = tensor_model_parallel_all_gather(m.bias, dim=0)
+                bias = bias[:m.org_vocab_size].contiguous()
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(bias.cpu(), requires_grad=False)
+            
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            weight_tensor = weight_tensor[:m.org_vocab_size,:].contiguous()
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(weight_tensor.cpu(), requires_grad=False)
+                print(f"merged: {name}, shape={m.weight.shape}")
+        
+        elif isinstance(m, QKVParallelLinear):
+            # weight shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp, hidden_size]
+            # bias shape: [total_num_head * head_size // tp + 2 * total_num_head * head_size // tp]
+            if self.parallel_config.world_size > 1:
+                total_q_hidden_size = m.total_num_heads * m.head_size
+                partial_q_hidden_size = m.num_heads * m.head_size
+                total_kv_hidden_size = m.total_num_kv_heads * m.head_size
+                partial_kv_hidden_size = m.num_kv_heads * m.head_size
+
+                if m.bias is not None:
+                    bias_tenosr = m.bias.new_zeros(total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size)
+                    q_bias = bias_tenosr[:total_q_hidden_size][self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
+                    k_bias = bias_tenosr[total_q_hidden_size:total_q_hidden_size+total_kv_hidden_size]\
+                                [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+                    v_bias = bias_tenosr[total_q_hidden_size+total_kv_hidden_size:]\
+                                [self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+                    
+                    q_bias[:] = m.bias[:partial_q_hidden_size]
+                    k_bias[:] = m.bias[partial_q_hidden_size:partial_q_hidden_size+partial_kv_hidden_size]
+                    v_bias[:] = m.bias[partial_q_hidden_size+partial_kv_hidden_size:]
+                    
+                    bias_tensor = tensor_model_parallel_all_reduce(bias_tenosr)
+                    if self.is_driver_worker:
+                        m.bias = torch.nn.Parameter(bias_tensor.cpu(), requires_grad=False)
+                
+                weight_tensor = m.weight.new_zeros(
+                    total_q_hidden_size + total_kv_hidden_size * 2, m.hidden_size
+                )
+
+                q_tensor = weight_tensor[:total_q_hidden_size, :]
+                q_tensor = q_tensor[self.rank * partial_q_hidden_size : (self.rank + 1) * partial_q_hidden_size]
+
+                k_tensor = weight_tensor[total_q_hidden_size : total_q_hidden_size + total_kv_hidden_size]
+                k_tensor = k_tensor[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+
+                v_tensor = weight_tensor[total_q_hidden_size + total_kv_hidden_size :]
+                v_tensor = v_tensor[self.rank * partial_kv_hidden_size : (self.rank + 1) * partial_kv_hidden_size]
+
+                q_tensor[:, :] = m.weight[: partial_q_hidden_size, :]
+                k_tensor[:, :] = m.weight[partial_q_hidden_size : partial_q_hidden_size + partial_kv_hidden_size, :]
+                v_tensor[:, :] = m.weight[partial_q_hidden_size + partial_kv_hidden_size : , :]
+
+                weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
+            else:
+                weight_tensor = m.weight
+                if m.bias is not None and self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+                    
+            int8_weight, weight_scales = w8a16_quantization(weight_tensor)
+            
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
+                m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
+                print(f"Quantized: {name}")
+
+        elif isinstance(m, MergedColumnParallelLinear):
+            # weight shape: [intermediate_size // tp * 2, hidden_size]
+            # bias shape: [intermediate_size // tp * 2]
+            if self.parallel_config.world_size > 1:
+                output_sizes = m.output_sizes
+                output_size = sum(output_sizes)
+                partial_output_sizes = [
+                    i // self.parallel_config.world_size for i in output_sizes
+                ]
+                
+                if m.bias is not None:
+                    index_start = 0
+                    partial_index_start = 0
+                    bias_tenosr = m.bias.new_zeros(output_size)
+                    for i in range(len(output_sizes)):
+                        index_out = index_start + output_sizes[i]
+                        sub_bias_tensor = bias_tenosr[index_start:index_out]
+                        partial_size = partial_output_sizes[i]
+                        sub_bias_tensor[self.rank * partial_size:(self.rank+1) * partial_size] = m.bias[partial_index_start:partial_index_start+partial_size]
+                        
+                        index_start += output_sizes[i]
+                        partial_index_start += partial_size
+                    bias_tenosr = tensor_model_parallel_all_reduce(bias_tenosr)
+                    if self.is_driver_worker:
+                        m.bias = torch.nn.Parameter(bias_tenosr, requires_grad=False)
+
+                weight_tensor = m.weight.new_zeros(output_size, m.input_size)
+
+                index_start = 0
+                partial_index_start = 0
+                for i in range(len(output_sizes)):
+                    index_out = index_start + output_sizes[i]
+                    sub_weight_tensor = weight_tensor[index_start:index_out]
+                    partial_size = partial_output_sizes[i]
+                    sub_weight_tensor[self.rank * partial_size : (self.rank + 1) * partial_size] = m.weight[partial_index_start : partial_index_start + partial_size]
+
+                    index_start += output_sizes[i]
+                    partial_index_start += partial_size
+                weight_tensor = tensor_model_parallel_all_reduce(weight_tensor)
+            else:
+                weight_tensor = m.weight
+                if m.bias is not None and self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+                
+            int8_weight, weight_scales = w8a16_quantization(weight_tensor)
+            
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
+                m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
+                print(f"Quantized: {name}")
+        
+        elif isinstance(m, ColumnParallelLinear):
+            # weight shape: [some_dim // tp, hidden_size] // for this Linear, some_dim mostly is hidden_size * 4
+            # bias shape: [some_dim // tp]
+            if m.bias is not None:
+                bias_tenosr = tensor_model_parallel_all_gather(m.bias, dim=0)
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(bias_tenosr.cpu(), requires_grad=False)
+            
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=0)
+            
+            int8_weight, weight_scales = w8a16_quantization(weight_tensor)
+            
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
+                m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
+                print(f"Quantized: {name}")
+        
+        elif isinstance(m, RowParallelLinear):
+            # weight shape: [hidden_size, some_dim // tp] // for this Linear, some_dim mostly is hidden_size * 4 or intermediate_size
+            # bias shape: [hidden_size]
+            if m.bias is not None:
+                bias_tensor = tensor_model_parallel_all_gather(m.bias, dim=-1)
+                if self.is_driver_worker:
+                    m.bias = torch.nn.Parameter(m.bias.cpu(), requires_grad=False)
+            
+            weight_tensor = tensor_model_parallel_all_gather(m.weight, dim=-1)
+            int8_weight, weight_scales = w8a16_quantization(weight_tensor)
+            
+            if self.is_driver_worker:
+                m.weight = torch.nn.Parameter(int8_weight.cpu(), requires_grad=False)
+                m.scales = torch.nn.Parameter(weight_scales.cpu(), requires_grad=False)
+                print(f"Quantized: {name}")
+        else:
+            pass
+        
+        torch.cuda.empty_cache()
+    
+    # save weights
+    if self.is_driver_worker:
+        from safetensors.torch import save_file
+
+        tensors = {}
+        saved = False
+        count = 0
+        size_in_bytes = 0
+        
+        for name, weight in model.named_parameters():
+            if "lm_head" in name and model.config.tie_word_embeddings:
+                continue
+            size_in_bytes += weight.numel() * weight.element_size()
+            tensors[name] = weight
+            saved = False
+            if gb_per_file is not None and size_in_bytes >= gb_per_file * 1024 * 1024 * 1024:
+                weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))
+                save_file(tensors, weight_path)
+                print(f"The quantified weights were successfully saved in {weight_path}.")
+                tensors.clear()
+                saved = True
+                count += 1
+                size_in_bytes = 0
+        
+        if not saved: 
+            weight_path = os.path.join(save_path, "model_{}.safetensors".format(str(count).zfill(6)))   
+            save_file(tensors, weight_path)
+            print(f"The quantified weights were successfully saved in {weight_path}.")

ixformer_sdk/contrib/vllm_flash_attn/__init__.py

@@ -0,0 +1,3 @@
+__version__ = "2.6.1"
+
+from .flash_attn_interface import *

ixformer_sdk/core/config.py

@@ -0,0 +1,184 @@
+import os
+from typing import Callable, Optional
+
+# =========================================================
+# Utils
+# =========================================================
+
+
+def number_type(scalar_type):
+    def wrap(val: Optional[str]):
+        if val is None:
+            return None
+
+        return scalar_type(val)
+
+    return wrap
+
+
+def bool_type(val: Optional[str]):
+    if val is None:
+        return False
+
+    if isinstance(val, str):
+        return val.lower() in ["1", "t", "true"]
+
+    if isinstance(val, int):
+        return val != 0
+
+    raise RuntimeError(f"Invalid bool type, got {type(val), val}")
+
+
+def list_type(scalar_type=str):
+    def wrap(val: Optional[str]):
+        if val is None:
+            return []
+
+        if not isinstance(val, str):
+            raise RuntimeError(
+                f"list_type: Got invalid type, expect str, but got {val}."
+            )
+
+        return [scalar_type(v) for v in val.split(",")]
+
+    return wrap
+
+
+def Field(
+    name: str,
+    static: bool = True,
+    type: Callable = str,
+    choices: Optional[list] = None,
+    help: Optional[str] = None,
+    **kwargs,
+):
+    """
+    Define environment variable field
+
+    Example:
+    Static mode:
+        # define
+        ENABLE_XX = Field("ENABLE_XX", type=bool, help="ENABLE_XX")
+
+        # use
+        config.ENABLE_XX
+
+    Dynamic mode:
+        # Please use lowercase naming to differentiate it with static mode.
+
+        # define
+        enable_cc = Field("ENABLE_CC", type=bool, static=False, help="enable_cc")
+
+        # use
+        config.enable_cc()
+
+    Set default value:
+        # define
+        ENABLE_TT = Field("ENABLE_TT", type=bool, default=False, help="ENABLE_TT")
+
+        # use
+        config.ENABLE_TT
+
+    Use list:
+        # define
+        CUDA_VISIBLE_DEVICES = Field("CUDA_VISIBLE_DEVICES", type=list_type(int), help="CUDA_VISIBLE_DEVICES")
+
+        # use
+        # the CUDA_VISIBLE_DEVICES is parsed to list, and it's value is int type.
+        for device_id in CUDA_VISIBLE_DEVICES:
+            ...
+
+    """
+
+    if type == bool:
+        type = bool_type
+
+    elif type in [list, tuple]:
+        type = list_type(scalar_type=str)
+
+    elif type in [int, float]:
+        type = number_type(type)
+
+    if static:
+        env_val = type(os.environ.get(name, **kwargs))
+        if choices is not None and env_val is not None and env_val not in choices:
+            raise RuntimeError(
+                f"Got invalid value, expect {choices}, but got {env_val}."
+            )
+        return env_val
+
+    def _get():
+        env_val = type(os.environ.get(name, **kwargs))
+        if choices is not None and env_val is not None and env_val not in choices:
+            raise RuntimeError(
+                f"Got invalid value, expect {choices}, but got {env_val}."
+            )
+        return env_val
+
+    return _get
+
+
+# =========================================================
+# Functions Config
+# =========================================================
+
+IXFORMER_GEMV_THRESHOLD = Field(
+    "IXFORMER_GEMV_THRESHOLD",
+    type=int,
+    default=1,
+    help="Set the threshold for using gemv.",
+)
+
+
+# =========================================================
+# Distributed Config
+# =========================================================
+
+IXFORMER_COMM_SHM_SIZE = Field(
+    "IXFORMER_COMM_SHM_SIZE",
+    type=int,
+    default=None,
+    help="set shared memory size of ipc comm.",
+)
+
+IXFORMER_ENABLE_OVERLAP_COMM = Field(
+    "IXFORMER_ENABLE_OVERLAP_COMM",
+    type=bool,
+    default=False,
+    help="enable overlap communcation and compute.",
+)
+
+IXFORMER_OVERLAP_GEMM_METHOD = Field(
+    "IXFORMER_OVERLAP_GEMM_METHOD",
+    type=int,
+    default=None,
+    choices=[0, 1],
+    help="set gemm backend, 0: ixinfer, 1: cublas.",
+)
+
+IXFORMER_OVERLAP_CHUNKS = Field(
+    "IXFORMER_OVERLAP_CHUNKS", type=int, default=2, help="set split chunks."
+)
+
+IXFORMER_OVERLAP_SPLIT_RATIO = Field(
+    "IXFORMER_OVERLAP_SPLIT_RATIO",
+    type=float,
+    default=None,
+    help="set split chunks ratio.",
+)
+
+IXFORMER_PAGED_ATTENTION_ALGO = Field(
+    "IXFORMER_PAGED_ATTENTION_ALGO",
+    type=str,
+    default="ixinfer",
+    choices=["ixinfer", "ixformer"],
+    help="set paged attention algo.",
+)
+
+IXFORMER_UNPAD_ATTENTION_ALGO = Field(
+    "IXFORMER_UNPAD_ATTENTION_ALGO",
+    type=str,
+    default="ixinfer",
+    choices=["ixinfer", "ixinfer-ex"],
+    help="set enpad attention algo.",
+)

ixformer_sdk/core/dispatcher.py

@@ -0,0 +1,20 @@
+class Dispatcher(object):
+    """
+    create object by dispatcher to reuse object.
+    """
+
+    _dispatcher = dict()
+
+    @classmethod
+    def dispatcher(cls, *args, **kwargs):
+        key = cls.dispatcher_key(*args, **kwargs)
+        obj = cls._dispatcher.get(key, None)
+        if obj is None:
+            obj = cls(*args, **kwargs)
+            cls._dispatcher[key] = obj
+
+        return obj
+
+    @classmethod
+    def dispatcher_key(cls, *args, **kwargs):
+        raise NotImplementedError()

ixformer_sdk/core/multi_level_cache.py

@@ -0,0 +1,54 @@
+class MultiLevelCache(object):
+    def __init__(self):
+        self._l1_key = None
+        self._l1_value = None
+
+        self._l2_size = 3
+        self._l2 = [(None, None) for _ in range(self._l2_size)]
+        self._l2_ptr = 0
+
+        self._l3 = dict()
+
+    def set(self, key, value):
+        self._l1_key = key
+        self._l1_value = value
+
+        self._l2[self._l2_ptr] = (key, value)
+        self._l2_ptr = (self._l2_ptr + 1) % 3  # l2_size: 3
+
+        self._l3[key] = value
+
+    def get(self, key, *args):
+        if key == self._l1_key:
+            return self._l1_value
+
+        l2 = self._l2
+        if key == l2[0][0]:
+            return l2[0][1]
+
+        if key == l2[1][0]:
+            return l2[1][1]
+
+        if key == l2[2][0]:
+            return l2[2][1]
+
+        return self._l3.get(key, *args)
+
+    def containe(self, key):
+        return key in self._l3
+
+    def __getitem__(self, item):
+        return self.get(item)
+
+    def __setitem__(self, key, value):
+        self.set(key, value)
+
+    def __contains__(self, item):
+        if item == self._l1_key:
+            return True
+
+        l2 = self._l2
+        if item == l2[0][0] or item == l2[1][0] or item == l2[2][0]:
+            return True
+
+        return item in self._l3

ixformer_sdk/core/operator_autotuning.py

@@ -0,0 +1,237 @@
+import abc
+import bisect
+import functools
+import itertools
+import random
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+
+import ixformer.distributed as ixfd
+from ixformer.utils.benchmark.cuda_benchmark import Functor, cuda_benchmark
+
+
+def sync_ranks_metric(value, group=None):
+    if not isinstance(value, (torch.Tensor, int, float)):
+        raise RuntimeError(
+            f"Invalid metric value, expect `Tensor`, `int`, or `float` type, but got {value}."
+        )
+
+    if torch.is_tensor(value):
+        value = value.to("cuda")
+    else:
+        value = torch.tensor([value], dtype=torch.float, device="cuda")
+
+    dist.broadcast(value, src=0, group=group)
+    return value.cpu().item()
+
+
+class AutotuningFinder(object):
+    def freeze(self):
+        pass
+
+    @abc.abstractmethod
+    def get(self, key) -> Callable:
+        pass
+
+    @abc.abstractmethod
+    def set(self, *args, **kwargs):
+        pass
+
+
+class BasedKeyFinder(AutotuningFinder):
+    def __init__(self):
+        self._key_to_value: Dict[Any, Callable] = dict()
+
+    def get(self, key, **kwargs) -> Callable:
+        if "default" in kwargs:
+            return self._key_to_value.get(key, kwargs["default"])
+        return self._key_to_value[key]
+
+    def set(self, key, value):
+        self._key_to_value[key] = value
+
+    def containe(self, key):
+        return key in self._key_to_value
+
+
+class TreeNode:
+    def __init__(self):
+        self.nodes: List[Union[Any, TreeNode]] = list()
+        self.key_to_nodes: Dict[Any, TreeNode] = dict()
+
+    def add(self, key, value):
+        if isinstance(key, (tuple, list)):
+            if len(key) == 1:
+                self.insert_value(key[0], value)
+            else:
+                self.recurse_add_node(key, value)
+        else:
+            self.insert_value(key, value)
+
+    def insert_value(self, key, value):
+        self.nodes.append((key, value))
+        self.key_to_nodes[key] = value
+
+    def recurse_add_node(self, key, value):
+        if key[0] in self.key_to_nodes:
+            node = self.key_to_nodes[key[0]]
+        else:
+            node = TreeNode()
+            self.key_to_nodes[key[0]] = node
+            self.insert_value(key[0], node)
+
+        node.add(key[1:], value)
+
+    def sort(self):
+        self.nodes.sort(key=lambda x: x[0])
+        for _, node in self.nodes:
+            if isinstance(node, TreeNode):
+                node.sort()
+
+    def find(self, key):
+        is_list_key = isinstance(key, (tuple, list))
+        if not is_list_key:
+            key = (key,)
+
+        num_querys = len(key)
+        node = self
+        for key_idx in range(num_querys):
+            query_key = key[key_idx]
+            idx = bisect.bisect_left(node.nodes, (query_key,)) - 1
+            if idx <= 0:
+                node = node.nodes[0][1]
+            elif idx >= len(node.nodes):
+                node = node.nodes[-1][1]
+            else:
+                node = node.nodes[idx][1]
+
+        return node
+
+    def show(self, indent=0):
+        for k, node in self.nodes:
+            print(" " * indent, end="")
+            if isinstance(node, TreeNode):
+                print(f"key: {k}")
+                node.show(indent=indent + 4)
+            else:
+                print(f"key: {k}, node: {node}")
+
+
+class BasedRangeFinder(AutotuningFinder):
+    def __init__(self):
+        super().__init__()
+
+        self.tree = TreeNode()
+        self._found_cache: Dict[Any, Callable] = dict()
+
+    def freeze(self):
+        self.tree.sort()
+
+    def get(self, key) -> Callable:
+        value = self._found_cache.get(key, None)
+        if value is not None:
+            return value
+
+        value = self.tree.find(key)
+        self._found_cache[key] = value
+        return value
+
+    def set(self, key, value):
+        self.tree.add(key, value)
+
+
+class OperatorAutotuning(object):
+    def __init__(self, num_repeated=5, num_warmup=3, dist_barrier=False):
+        self.num_repeated = num_repeated
+        self.num_warmup = num_warmup
+        self.dist_barrier = dist_barrier
+
+    @abc.abstractmethod
+    def operators(self):
+        raise NotImplementedError()
+
+    def __call__(self, *args, **kwargs):
+        return self.exec_best_operator(args, kwargs)
+
+    @abc.abstractmethod
+    def exec_best_operator(self, args, kwargs):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def autotuning(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def perf_best_operator(self, *args, **kwargs) -> Callable:
+        best_operator = None
+        best_operator_time = float("inf")
+
+        for idx, operator in enumerate(self.operators()):
+            op_time = self.perf_operator_time(operator, *args, **kwargs)
+            if op_time < best_operator_time:
+                best_operator = operator
+                best_operator_time = op_time
+
+            # print(operator, op_time)
+
+        return best_operator
+
+    def perf_operator_time(self, op: Callable, *args, **kwargs) -> float:
+        fn = Functor(op, *args, **kwargs)
+        time = cuda_benchmark(fn, self.num_repeated, self.num_warmup, self.dist_barrier)
+        return time.gpu
+
+
+class OperatorRuntimeAutotuning(OperatorAutotuning):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.operator_finder = BasedKeyFinder()
+
+    @abc.abstractmethod
+    def get_operator_key(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def exec_best_operator(self, args, kwargs):
+        key = self.get_operator_key(*args, **kwargs)
+        operator = self.operator_finder.get(key, default=None)
+        if operator is None:
+            operator = self.perf_best_operator(*args, **kwargs)
+            self.operator_finder.set(key, operator)
+
+        return operator(*args, **kwargs)
+
+
+class OperatorPreBaseRangeAutotuning(OperatorAutotuning):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.operator_finder = BasedRangeFinder()
+        self._finished_autotuning = False
+
+    @abc.abstractmethod
+    def get_operator_key(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def generate_operator_inputs(self) -> Iterable[Tuple[Tuple, Dict]]:
+        raise NotImplementedError()
+
+    def exec_best_operator(self, args, kwargs):
+        if not self._finished_autotuning:
+            self.autotuning()
+
+        best_op = self.operator_finder.get(self.get_operator_key(*args, **kwargs))
+        return best_op(*args, **kwargs)
+
+    def autotuning(self):
+        for op_args, op_kwargs in self.generate_operator_inputs():
+            best_op = self.perf_best_operator(*op_args, **op_kwargs)
+            self.operator_finder.set(
+                self.get_operator_key(*op_args, **op_kwargs), best_op
+            )
+
+        self.operator_finder.freeze()
+        self._finished_autotuning = True
+        # self.operator_finder.tree.show()

ixformer_sdk/csrc/FindIXFORMER.cmake

@@ -0,0 +1,40 @@
+# use python to find ixformer libs and include
+if (CMAKE_VERSION VERSION_LESS 3.18)
+    set(DEV_MODULE Development)
+else()
+    set(DEV_MODULE Development.Module)
+endif()
+
+find_package(Python COMPONENTS Interpreter ${DEV_MODULE} REQUIRED)
+
+# find ixformer
+set(IXFORMER_FOUND FALSE)
+
+if("${Python_FOUND}" STREQUAL "TRUE")
+    execute_process(
+            COMMAND ${Python_EXECUTABLE} -c "import os, ixformer; print(os.path.dirname(ixformer.__file__))"
+            OUTPUT_VARIABLE IXFORMER_PYDIR
+            ERROR_VARIABLE PYTHON_ERROR
+            RESULT_VARIABLE PYTHON_RESULT
+            OUTPUT_STRIP_TRAILING_WHITESPACE
+            ERROR_STRIP_TRAILING_WHITESPACE
+    )
+    if ("${IXFORMER_PYDIR}" STREQUAL "")
+        message("-- Not found ixFormer")
+    else ()
+        message("-- Found ixFormer: ${IXFORMER_PYDIR}")
+        set(IXFORMER_FOUND TRUE)
+    endif ()
+endif()
+
+set(IXFORMER_COMM_LIBS "ixformer_comm")
+set(IXFORMER_KERNEL_LIBS "ixformer_kernels")
+set(IXFORMER_LIBS "${IXFORMER_COMM_LIBS} ${IXFORMER_KERNEL_LIBS}")
+set(IXFORMER_INCLUDE "")
+set(IXFORMER_DIR "")
+
+if("${IXFORMER_FOUND}" STREQUAL "TRUE")
+    set(IXFORMER_INCLUDE "${IXFORMER_PYDIR}/csrc/include")
+    set(IXFORMER_DIR "${IXFORMER_PYDIR}")
+    message("-- ixFormer LIBS: ${IXFORMER_LIBS}, INCLUDE: ${IXFORMER_INCLUDE}")
+endif ()

ixformer_sdk/csrc/include/ixformer/comm/ccl.h

@@ -0,0 +1,357 @@
+#pragma once
+
+#include "core/op_algo.h"
+#include "nccl.h"
+
+namespace ixformer::comm {
+
+const uint8_t MAX_TENSOR_NDIM = 8;
+constexpr size_t DEFAULT_SHM_SIZE = 16 * 1024 * 1024 * sizeof(float);
+
+struct Comm;
+typedef Comm *Comm_t;
+
+
+struct TensorDesc {
+    void *data_ptr;
+    ncclDataType_t dtype;
+    uint64_t numel;
+    uint8_t ndim;
+    int64_t shape[MAX_TENSOR_NDIM];
+    int64_t stride[MAX_TENSOR_NDIM];
+    bool contiguous;
+};
+
+
+/**
+ * @brief Generate unique communicator id.
+ *
+ * Generates an Id to be used in ncclCommInitRank. ncclGetUniqueId should be
+ * called once and the Id should be distributed to all ranks in the
+ * communicator before calling ncclCommInitRank.
+ *
+ * @param commId: the unique id of communicator, it is created in main rank, and broadcast other rank.
+ */
+void getUniqueId(ncclUniqueId *commId);
+
+/**
+ * @brief Serialize commId to string.
+ * @param commId: the unique id of communicator.
+ * @return: serialized string.
+ */
+std::string serializeUniqueId(const ncclUniqueId &commId);
+
+/**
+ * @brief Deserialize the string of commId.
+ * @param commIdStr: serialized string by serializeUniqueId.
+ * @param commId: output commId
+ */
+void deserializeUniqueId(const std::string &commIdStr, ncclUniqueId *commId);
+
+/**
+ * @brief Creates a new communicator (multi process version).
+ *
+ * Rank must be between 0 and nranks-1 and unique within a communicator clique.
+ * Each rank is associated to a CUDA device, which has to be set before calling ncclCommInitRank.
+ *
+ * It is important to ensure that the current process's CUDA device is set by cudaSetDevice before calling this function,
+ * otherwise, an exception will be thrown.
+ *
+ * @param comm: Communicator
+ * @param nranks: the number of ranks.
+ * @param commId: the unique id of communicator.
+ * @param rank: the rank of current process
+ * @param shm_size: Unlike NCCL, IxFormer communication relies on CUDA IPC for communication by shared memory.
+ *        If the shm_size is not provided, it will use the default value: DEFAULT_SHM_SIZE.
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void initRank(Comm_t *comm, int nranks, ncclUniqueId commId, int rank, size_t shm_size = DEFAULT_SHM_SIZE);
+
+/**
+ * @brief Finalize a communicator.
+ *
+ * ncclCommFinalize flushes all issued communications,
+ * and marks communicator state as ncclInProgress. The state will change to ncclSuccess
+ * when the communicator is globally quiescent and related resources are freed; then,
+ * calling ncclCommDestroy can locally free the rest of the resources (e.g. communicator
+ * itself) without blocking.
+ *
+ * @param comm: Communicator
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void destroy(Comm_t comm);
+
+void delete_comm_resuouces(Comm_t comm);
+
+/**
+ * @brief Whether is initiated.
+ * @param comm: Communicator
+ */
+bool isInitiated(Comm_t comm);
+
+/**
+ * @brief Gets ncclComm_t.
+ * @param comm: Communicator
+ */
+ncclComm_t getNcclComm(Comm_t comm);
+
+/**
+ * @brief Gets the number of ranks in the communicator clique
+ * @param comm: Communicator
+ */
+int getWorldSize(Comm_t comm);
+
+/**
+ * @brief Gets the number of nodes in the communicator clique
+ * @param comm: Communicator
+ */
+int getNumNodes(Comm_t comm);
+
+/**
+ * @brief Returns the user-ordered "rank" associated with the communicator.
+ * @param comm: Communicator
+ */
+int getRank(Comm_t comm);
+
+/**
+ * @brief Returns the cuda device number associated with the communicator.
+ * @param comm: Communicator
+ */
+int getDevice(Comm_t comm);
+
+/**
+ * @brief Gets shared memory size in the communicator clique
+ * @param comm: Communicator
+ */
+uint64_t getIpcShmSize(Comm_t comm);
+
+
+// ============================================================================
+// Collective communication operations
+//
+// Collective communication operations must be called separately for each
+// communicator in a communicator clique.
+//
+// They return when operations have been enqueued on the CUDA stream.
+//
+// Since they may perform inter-CPU synchronization, each call has to be done
+// from a different thread or process, or need to use Group Semantics (see
+// below).
+// ============================================================================
+
+/**
+ * @brief Barrier the member of the communicator
+ * @param comm: Communicator
+ * @param stream: CUDA Stream
+ */
+void barrier(Comm_t comm, cudaStream_t stream);
+
+/**
+ * @brief All-Gather
+ *
+ * Each device gathers sendcount values from other GPUs into senddata,
+ * receiving data from rank i at offset i*sendcount.
+ * Assumes recvcount is equal to nranks*sendcount, which means that recvdata
+ * should have a size of at least nranks*sendcount elements.
+ *
+ * In-place operations will happen if senddata == recvdata + rank * sendcount.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param recvdata: recv data
+ * @param sendcount: the number of send elements, it is not nbytes.
+ * @param dtype: data type
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void allGather(Comm_t comm, const void *senddata, void *recvdata, size_t sendcount, ncclDataType_t dtype,
+               cudaStream_t stream, AllGatherAlgo algo = AllGatherAlgo::kNone);
+
+/**
+ * @brief All-Reduce
+ *
+ * Reduces data arrays of length count in senddata using op operation, and
+ * leaves identical copies of result on each recvdata.
+ *
+ * In-place operation will happen if senddata == recvdata.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param recvdata: recv data
+ * @param count: the number of elements, it is not nbytes.
+ * @param dtype: data type
+ * @param op：Reduce type: ncclSum, ncclProd, ncclMin, ncclMax, ncclAvg
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void allReduce(Comm_t comm, const void *senddata, void *recvdata, size_t count, ncclDataType_t dtype,
+               ncclRedOp_t op, cudaStream_t stream, AllReduceAlgo algo = AllReduceAlgo::kNone);
+
+/**
+ * @brief Whether is supported non-contiguous tensors
+ * @param comm: Communicator
+ * @param dtype: data type
+ * @param shape: tensor shape
+ * @param ndim: the ndim of tensor
+ * @param numel: the number of tensor
+ * @param op: Reduce type: ncclSum, ncclProd, ncclMin, ncclMax, ncclAvg
+ * @return: supported
+ */
+bool allReduceStrideSupported(Comm_t comm, ncclDataType_t dtype, const int64_t *shape, int ndim, uint64_t numel, ncclRedOp_t op);
+
+/**
+ * @brief All-Reduce for non-contiguous tensors
+ * @param comm: Communicator
+ * @param senddata: send tensor
+ * @param recvdata: recv tensor
+ * @param stream: CUDA Stream
+ */
+void allReduceStride(Comm_t comm, const TensorDesc &senddata, TensorDesc &recvdata, cudaStream_t stream);
+
+/**
+ * @brief Send data from senddata to rank peer.
+ *
+ * Rank peer needs to call ncclRecv with the same datatype and the same count from this
+ * rank. This operation is blocking for the GPU. If multiple ncclSend and ncclRecv operations
+ * need to progress concurrently to complete.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param count: the number of send elements, it is not nbytes.
+ * @param dtype: data type
+ * @param peer: the destination rank
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void send(Comm_t comm, const void *senddata, size_t count, ncclDataType_t dtype, int peer, cudaStream_t stream,
+          SendAlgo algo = SendAlgo::kNone);
+
+/**
+ * @brief Receive data from rank peer into recvdata.
+ *
+ * Rank peer needs to call ncclSend with the same datatype and the same count to this
+ * rank. This operation is blocking for the GPU. If multiple ncclSend and ncclRecv operations
+ * need to progress concurrently to complete.
+ *
+ * @param comm: Communicator
+ * @param recvdata: recv data
+ * @param count: the number of recv elements, it is not nbytes.
+ * @param dtype: data type
+ * @param peer: source rank
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void recv(Comm_t comm, void *recvdata, size_t count, ncclDataType_t dtype, int peer, cudaStream_t stream,
+          RecvAlgo algo = RecvAlgo::kNone);
+
+/**
+ * @brief Reduces data arrays of length count in senddata into recvdata using op operation.
+ *
+ * Recvdata may be NULL on all calls except for root device.
+ * root is the rank (not the CUDA device) where data will reside after the
+ * operation is complete.
+ *
+ * In-place operation will happen if senddata == recvdata.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param recvdata: recv data
+ * @param count:  the number of elements, it is not nbytes.
+ * @param dtype: data type
+ * @param op: Reduce type: ncclSum, ncclProd, ncclMin, ncclMax, ncclAvg
+ * @param root: root rank
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void reduce(Comm_t comm, const void *senddata, void *recvdata, size_t count, ncclDataType_t dtype, ncclRedOp_t op,
+            int root, cudaStream_t stream, ReduceAlgo algo = ReduceAlgo::kNone);
+
+/**
+ * @brief Broadcast
+ *
+ * Copies count values from root to all other devices.
+ * root is the rank (not the CUDA device) where data resides before the
+ * operation is started.
+ *
+ * In-place operation will happen if senddata == recvdata.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param recvdata: recv data
+ * @param count: the number of elements, it is not nbytes.
+ * @param dtype: data type
+ * @param root: root rank
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void broadcast(Comm_t comm, const void *senddata, void *recvdata, size_t count, ncclDataType_t dtype, int root,
+               cudaStream_t stream, BroadcastAlgo algo = BroadcastAlgo::kNone);
+
+/**
+ *
+ * @brief Reduce-Scatter
+ *
+ * Reduces data in senddata using op operation and leaves reduced result
+ * scattered over the devices so that recvdata on rank i will contain the i-th
+ * block of the result.
+ * Assumes sendcount is equal to nranks*recvcount, which means that senddata
+ * should have a size of at least nranks*recvcount elements.
+ *
+ * In-place operations will happen if recvdata == senddata + rank * recvcount.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param recvdata: recv data
+ * @param recvcount: the number of recv elements, it is not nbytes.
+ * @param dtype: data type
+ * @param op：Reduce type: ncclSum, ncclProd, ncclMin, ncclMax, ncclAvg
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void reduceScatter(Comm_t comm, const void *senddata, void *recvdata,
+                   size_t recvcount, ncclDataType_t dtype, ncclRedOp_t op, cudaStream_t stream,
+                   ReduceScatterAlgo algo = ReduceScatterAlgo::kNone);
+
+/**
+ * @brief Send data from src_rank to dst_rank on src_rank process, recv data on dst_rank.
+ *
+ * @param comm: Communicator
+ * @param data: send data to dst rank if current rank is src_rank, recv data if current rank is dst_rank.
+ * @param count: the number of send/recv elements, it is not nbytes.
+ * @param dtype: data type
+ * @param src_rank: src rank
+ * @param dst_rank: dst rank
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void p2p(Comm_t comm, void *data, size_t count, ncclDataType_t dtype, int src_rank, int dst_rank, cudaStream_t stream,
+         SendAlgo algo = SendAlgo::kNone);
+
+/**
+ * @brief The all ranks of communicator send senddata to dst_rank.
+ *
+ * @param comm: Communicator
+ * @param senddata: send data
+ * @param recvdatas: recv datas，it is two-dim array, shape: [WorldSize, RecvDataPointer],
+ *        the first dim is host pointer，the second dim is GPU pointer,
+ *        it can be nullptr when current rank is not dst rank.
+ * @param sendcount: the number of send elements, it is not nbytes.
+ * @param dst_rank: dst rank
+ * @param dtype: data type
+ * @param stream: CUDA Stream
+ * @param algo: Algorithm
+ * @throw CommError: Throw CommError when an error is encountered.
+ */
+void gather(Comm_t comm, const void *senddata, void **recvdatas, size_t sendcount, int dst_rank, ncclDataType_t dtype,
+            cudaStream_t stream, GatherAlgo algo = GatherAlgo::kNone);
+
+
+}// namespace ixformer::comm

ixformer_sdk/csrc/include/ixformer/comm/core/error.h

@@ -0,0 +1,21 @@
+#pragma once
+
+#include <stdexcept>
+#include "status.h"
+
+namespace ixformer::comm {
+
+class CommError : public std::runtime_error {
+public:
+    template<class ERROR_STR>
+    CommError(CommStatus error, const ERROR_STR str) : error_{error}, std::runtime_error(str) {}
+
+    CommStatus status() {
+        return error_;
+    }
+
+private:
+    CommStatus error_;
+};
+
+}// namespace ixformer::comm

ixformer_sdk/csrc/include/ixformer/comm/core/op_algo.h

@@ -0,0 +1,80 @@
+#pragma once
+
+#include <string>
+
+namespace ixformer::comm {
+
+enum class AllGatherAlgo {
+    kNone,
+    kAuto,
+    kNCCL,
+    kNumAlgo
+};
+
+enum class AllReduceAlgo {
+    kNone,          // None
+    kAuto,          // 自动选择算法
+    kAllGatherSum,  // 针对小数据量的算法
+    kBroadcastSum,  // 针对小数据量的算法
+    kRing,          // Ring AllReduce
+    kQuant,         // 对通讯算法进行量化，默认使用 kQuantL1
+    kQuantL1,       // 对通讯算法进行量化，优先使用量化算法以及最大保留精度，在部分 Size 性能不佳时，退化为 Auto 算法
+    kQuantL2,       // 对通讯算法进行量化，优先使用量化算法以及最大化速度，在部分 Size 性能不佳时，退化为 Auto 算法
+    kQuantL1AllSize,// 对所有的 Size 都使用量化算法
+    kQuantL2AllSize,// 对所有的 Size 都使用量化算法
+    kNCCL,          // 使用 NCCL
+    kStride,        // 输入或输出的 Tensor 不是连续的
+    kNumAlgo
+};
+
+
+enum class BroadcastAlgo {
+    kNone,
+    kAuto,
+    kNCCL,
+    kNumAlgo
+};
+
+enum class GatherAlgo {
+    kNone,
+    kAuto,
+    kNCCL,
+    kNumAlgo
+};
+
+enum class SendAlgo {
+    kNone,
+    kAuto,
+    kNCCL,
+    kNumAlgo
+};
+
+typedef SendAlgo RecvAlgo;
+
+enum class ReduceAlgo {
+    kNone,
+    kAuto,
+    kNCCL,
+    kNumAlgo
+};
+
+enum class ReduceScatterAlgo {
+    kNone,
+    kAuto,
+    kNCCL,
+    kNumAlgo
+};
+
+
+std::string to_string(AllGatherAlgo algo);
+std::string to_string(AllReduceAlgo algo);
+std::string to_string(BroadcastAlgo algo);
+std::string to_string(GatherAlgo algo);
+std::string to_string(SendAlgo algo);
+std::string to_string(ReduceAlgo algo);
+std::string to_string(ReduceScatterAlgo algo);
+
+template<typename Algo>
+Algo get_algo_from_str(const std::string &name);
+
+}// namespace ixformer::comm

ixformer_sdk/csrc/include/ixformer/comm/core/status.h

@@ -0,0 +1,22 @@
+#pragma once
+#include "comm/core/common.h"
+
+
+namespace ixformer::comm {
+
+enum CommStatus {
+    commSuccess,
+    commFail,
+    commCudaError,
+    commNcclError,
+    commInvalidArgument,
+    commUnsupported,
+    commInternalError,
+    commInvalidComm// maybe comm is nullptr
+};
+
+
+std::string to_string(CommStatus status);
+
+
+}// namespace ixformer::comm

ixformer_sdk/csrc/include/ixformer/kernels/error.h

@@ -0,0 +1,22 @@
+#pragma once
+
+#include <stdexcept>
+
+#include "status.h"
+
+namespace ixformer::kernels {
+
+class KernelError : public std::runtime_error {
+public:
+    template<class ERROR_STR>
+    KernelError(KernelStatus error, const ERROR_STR str) : error_{error}, std::runtime_error(str) {}
+
+    KernelStatus status() {
+        return error_;
+    }
+
+private:
+    KernelStatus error_;
+};
+
+}// namespace ixformer::kernels

ixformer_sdk/csrc/include/ixformer/kernels/status.h

@@ -0,0 +1,20 @@
+#pragma once
+
+#include <string>
+
+namespace ixformer::kernels {
+
+enum KernelStatus {
+    kernelSuccess,
+    kernelFail,
+    kernelCudaError,
+    kernelInvalidArgument,
+    kernelCuinferError,
+    kernelUnsupported,
+};
+
+
+std::string to_string(KernelStatus status);
+
+
+}// namespace ixformer::kernels

ixformer_sdk/csrc/include/ixformer/kernels/tensor.h

@@ -0,0 +1,92 @@
+#pragma once
+
+#include <string>
+
+namespace ixformer::kernels {
+
+const uint8_t MAX_TENSOR_NDIM = 8;
+
+// align with at::ScalarType
+enum DType {
+    Byte = 0,
+    Char = 1,
+    Short = 2,
+    Int = 3,
+    Long = 4,
+    Half = 5,
+    Float = 6,
+    Double = 7,
+    ComplexHalf = 8,
+    ComplexFloat = 9,
+    ComplexDoubl = 10,
+    Bool = 11,
+    QInt8 = 12,
+    QUInt8 = 13,
+    QInt32 = 14,
+    BFloat16 = 15,
+    QUInt4x2 = 16,
+    QUInt2x4 = 17,
+    Bits1x8 = 18,
+    Bits2x4 = 19,
+    Bits4x2 = 20,
+    Bits8 = 21,
+    Bits16 = 22,
+    Float8_e5m2 = 23,
+    Float8_e4m3fn = 24,
+    Undefined = 25,
+    NumOptions = 26
+};
+
+struct TensorDesc {
+
+public:
+    // delete default constructor
+    TensorDesc() = delete;
+    // All information must be (should be) prepared when constructing a TensorDesc object.
+    TensorDesc(DType scalar_type, void *data_ptr, int64_t numel, int64_t dim, const int64_t *size, const int64_t *stride, bool is_contiguous, bool is_cuda)
+        : dtype(scalar_type), ptr(data_ptr), nnumel(numel), ndim(dim), sizes(size), strides(stride), contiguous(is_contiguous), cuda(is_cuda) {}
+
+    inline DType scalar_type() const {
+        return dtype;
+    }
+
+    inline void *data_ptr() const {
+        return ptr;
+    }
+
+    inline int64_t numel() const {
+        return nnumel;
+    }
+
+    inline int64_t dim() const {
+        return ndim;
+    }
+
+    inline int64_t size(int64_t dim) const {
+        return dim < 0 ? sizes[ndim - dim] : sizes[dim];
+    }
+
+    inline int64_t stride(int64_t dim) const {
+        return dim < 0 ? strides[ndim - dim] : strides[dim];
+    }
+
+    inline bool is_contiguous() const {
+        return contiguous;
+    }
+
+    inline bool is_cuda() const {
+        return cuda;
+    }
+
+private:
+    void *ptr{nullptr};
+    DType dtype;
+    int64_t nnumel{0};
+    int64_t ndim{0};
+    const int64_t *sizes{nullptr};
+    const int64_t *strides{nullptr};
+    bool contiguous{false};
+    bool cuda{false};
+};
+
+}// namespace ixformer::kernels

ixformer_sdk/distributed/__init__.py

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

ixformer_sdk/distributed/_distributed.py

@@ -0,0 +1,481 @@
+import warnings
+from collections import defaultdict
+from typing import List, Optional, Tuple
+
+import torch
+import torch.distributed as dist
+import torch.distributed.distributed_c10d as c10d
+from ixformer._C import _distributed as cdist
+from ixformer._C._distributed import comm
+from ixformer._C._distributed.comm import (
+    AllGatherAlgo,
+    AllReduceAlgo,
+    BroadcastAlgo,
+    ReduceAlgo,
+    ReduceOp,
+    ReduceScatterAlgo,
+    SendAlgo,
+)
+from ixformer.core.multi_level_cache import MultiLevelCache
+from torch import Tensor
+from torch.distributed import ProcessGroup
+
+from ixformer.core import config
+
+IxformerCommType = int
+RecvAlgo = SendAlgo
+
+_GROUP_TO_IXFC_COMM_CACHE = MultiLevelCache()
+_IXFC_COMM_TO_GROUP_CACHE = MultiLevelCache()
+
+
+def get_store(group: dist.ProcessGroup = None) -> dist.Store:
+    if group is None:
+        group = c10d._get_default_group()
+
+    return c10d._pg_map[group][1]
+
+
+class StoreWrapper(cdist.comm.C10dStoreWrapper):
+    _GROUP_COUNT = defaultdict(dict)
+
+    def __init__(self, group: ProcessGroup):
+        super().__init__()
+        self.store = get_store()
+
+        ranks = dist.get_process_group_ranks(group)
+
+        group_key = "_".join([str(r) for r in ranks])
+        if group not in self._GROUP_COUNT[group_key]:
+            self._GROUP_COUNT[group_key][group] = len(self._GROUP_COUNT[group_key])
+        group_count = self._GROUP_COUNT[group_key][group]
+
+        self.prefix = f"gid_{group_count}_" + group_key
+
+    def _gen_unique_key(self, key):
+        return f"{self.prefix}_{key}"
+
+    def set(self, key: str, value: str):
+        key = self._gen_unique_key(key)
+        self.store.set(key, value)
+
+    def get(self, key: str) -> str:
+        key = self._gen_unique_key(key)
+        self.store.wait([key])
+        return self.store.get(key).decode("utf8")
+
+
+def init_comm_with_store(group=None, shmsize: int = None):
+    if group is None:
+        group = c10d._get_default_group()
+
+    world_size = dist.get_world_size(group=group)
+    rank = dist.get_group_rank(group=group, global_rank=dist.get_rank())
+
+    if shmsize is None:
+        shmsize = config.IXFORMER_COMM_SHM_SIZE
+
+    store_wrapper = StoreWrapper(group=group)
+    ixfc_comm = cdist.comm.init_communicator_by_store(
+        store=store_wrapper, world_size=world_size, rank=rank, max_shm_mem_size=shmsize
+    )
+
+    _GROUP_TO_IXFC_COMM_CACHE.set(group, ixfc_comm)
+    _IXFC_COMM_TO_GROUP_CACHE.set(ixfc_comm, group)
+    return ixfc_comm
+
+
+_sub_store = None
+
+
+def create_nccl_unique_id(addr: str, port: str, world_size: int, rank: int):
+    global _sub_store
+    _sub_store = dist.TCPStore(
+        host_name=addr, port=int(port), world_size=world_size, is_master=rank == 0
+    )
+    store_key = "ncclUniqueId"
+    if rank == 0:
+        commid = cdist.comm.create_nccl_unique_id()
+        _sub_store.set(store_key, commid)
+    else:
+        _sub_store.wait([store_key])
+        commid = _sub_store.get(store_key).decode("utf8")
+
+    return commid
+
+
+def init_comm_with_eth(
+    addr: str, port: str, world_size: int, rank: int, shmsize: int = None
+):
+    commid = create_nccl_unique_id(addr, port, world_size=world_size, rank=rank)
+    return cdist.comm.init_communicator_by_nccl_id(commid, world_size, rank, shmsize)
+
+
+def _check_group(group: Optional[ProcessGroup] = None):
+    if group is None:
+        group = c10d._get_default_group()
+
+    if isinstance(group, ProcessGroup):
+        ixfc_comm = _GROUP_TO_IXFC_COMM_CACHE.get(group, None)
+        if ixfc_comm is None:
+            return init_comm_with_store(group)
+        return ixfc_comm
+
+    return group
+
+
+def get_comm_group_stream(group: Optional[ProcessGroup] = None):
+    group = _check_group(group)
+    return comm.get_comm_group_stream(group)
+
+
+def set_comm_group_stream(stream: int, group: Optional[ProcessGroup] = None):
+    group = _check_group(group)
+    return comm.set_comm_group_stream(group, stream)
+
+
+def get_group_rank(group: Optional[ProcessGroup], global_rank) -> int:
+    """将 global rank 映射到 group 中的相对 rank"""
+    if isinstance(group, IxformerCommType):
+        _pg = _IXFC_COMM_TO_GROUP_CACHE.get(group, None)
+        if _pg is None:
+            return global_rank
+        else:
+            group = _IXFC_COMM_TO_GROUP_CACHE.get(group)
+
+    if group is None:
+        group = c10d._get_default_group()
+
+    return dist.get_group_rank(group, global_rank)
+
+
+def get_global_rank(group: Optional[ProcessGroup], group_rank: int) -> int:
+    """将一个 group rank 映射到 global rank"""
+    if group is None:
+        group = c10d._get_default_group()
+    return c10d.get_global_rank(group, group_rank)
+
+
+def get_process_group_ranks(group: Optional[ProcessGroup] = None) -> List[int]:
+    """获取 Group 的 global ranks"""
+    if group is None:
+        group = c10d._get_default_group()
+    return c10d.get_process_group_ranks(group)
+
+
+def new_group(ranks: List[int] = None, shmsize=None, *args, **kwargs):
+    """通过 global ranks 去创建一个通讯组"""
+    group = c10d.new_group(ranks, *args, **kwargs)
+
+    if ranks is None:
+        ranks = dist.get_process_group_ranks(group)
+
+    if get_rank() in ranks:
+        init_comm_with_store(group=group, shmsize=shmsize)
+    return group
+
+
+def new_subgroups_by_enumeration(
+    ranks_per_subgroup_list, shmsize=None, *args, **kwargs
+) -> Tuple[ProcessGroup, List[ProcessGroup]]:
+    """
+    通过一组 global ranks 去创建通讯组
+
+    :param ranks_per_subgroup_list: global ranks
+    :return: 返回当前 rank 所在的通讯组 和 新的 subgroups
+    """
+    self_group, other_group = c10d.new_subgroups_by_enumeration(
+        ranks_per_subgroup_list, *args, **kwargs
+    )
+    init_comm_with_store(self_group, shmsize=shmsize)
+    return self_group, other_group
+
+
+def destroy_process_group(group: Optional[ProcessGroup] = None):
+    """销毁 Group"""
+    if group is None:
+        group = c10d._get_default_group()
+    ixfc_comm = _GROUP_TO_IXFC_COMM_CACHE.get(group, None)
+
+    if ixfc_comm is None:
+        dist.destroy_process_group(group)
+    else:
+        comm.destroy(ixfc_comm)
+        dist.destroy_process_group(group)
+
+
+def get_rank(group: Optional[ProcessGroup] = None) -> int:
+    """获取当前进程的 Rank，如果 group 是 null，那么返回的是 Global Rank, 否则返回的相对的 Rank，即在当前组中的 rank"""
+    return c10d.get_rank(group)
+
+
+def get_world_size(group: Optional[ProcessGroup] = None) -> int:
+    """获取 Group 中的成员大小"""
+    return c10d.get_world_size(group)
+
+
+def barrier(group: Optional[ProcessGroup] = None, use_comm_stream: bool = False):
+    """同步 Group 中的 rank"""
+    group = _check_group(group)
+    comm.barrier(group, use_comm_stream)
+
+
+def isend(
+    tensor: Tensor,
+    dst: int,
+    group: Optional[ProcessGroup] = None,
+    use_comm_stream: bool = False,
+):
+    dst = get_group_rank(group, dst)
+    group = _check_group(group)
+    return comm.send(group, tensor, dst, use_comm_stream, SendAlgo.kNone)
+
+
+def send(*args, **kwargs):
+    warnings.warn("not support sync mode, as async to call.")
+    return isend(*args, **kwargs)
+
+
+def irecv(
+    tensor: torch.Tensor,
+    src: int,
+    group: Optional[ProcessGroup] = None,
+    use_comm_stream: bool = False,
+):
+    src = get_group_rank(group, src)
+    group = _check_group(group)
+    return comm.recv(group, tensor, src, use_comm_stream, SendAlgo.kNone)
+
+
+def recv(*args, **kwargs):
+    warnings.warn("not support sync mode, as async to call.")
+    return irecv(*args, **kwargs)
+
+
+def point_to_point(
+    tensor: Tensor,
+    src: int,
+    dst: int,
+    group: Optional[ProcessGroup] = None,
+    use_comm_stream: bool = False,
+):
+    """在 src rank 发送 tensor，在 dst_rank 上接收数据到 tensor 中"""
+    src = get_group_rank(group, src)
+    dst = get_group_rank(group, dst)
+    group = _check_group(group)
+    return comm.p2p(group, tensor, src, dst, use_comm_stream)
+
+
+def reduce(
+    tensor,
+    root: int,
+    op=ReduceOp.SUM,
+    group: Optional[ProcessGroup] = None,
+    async_op=False,
+    out: Tensor = None,
+    use_comm_stream: bool = False,
+):
+    """
+    Example:
+        ixf_tensor = torch.tensor([1], device="cuda")
+        ixfd.reduce(ixf_tensor, 1, async_op=True)
+        print("rank {rank}:", ixf_tensor)
+
+        # output
+        rank 0:  tensor([1], device='cuda:0')
+        rank 1:  tensor([4], device='cuda:1')
+        rank 2:  tensor([1], device='cuda:2')
+        rank 3:  tensor([1], device='cuda:3')
+    """
+
+    if not async_op:
+        raise RuntimeError("Not support sync operation now.")
+
+    if out is None:
+        out = tensor
+
+    root = get_group_rank(group, root)
+    group = _check_group(group)
+    return comm.reduce(group, tensor, out, op, root, use_comm_stream, ReduceAlgo.kNone)
+
+
+def broadcast(
+    tensor: Tensor,
+    src: int,
+    group: Optional[ProcessGroup] = None,
+    async_op=False,
+    out: Tensor = None,
+    use_comm_stream: bool = False,
+):
+    """
+    Example:
+        ixf_tensor = torch.tensor([rank], device="cuda")
+        ixfd.broadcast(ixf_tensor, 1, async_op=True)
+        print("rank {rank}: ", ixf_tensor)
+
+        # output
+        rank 0:  tensor([1], device='cuda:0')
+        rank 1:  tensor([1], device='cuda:1')
+        rank 2:  tensor([1], device='cuda:2')
+        rank 3:  tensor([1], device='cuda:3')
+    """
+    if not async_op:
+        raise RuntimeError("Not support sync operation now.")
+
+    if out is None:
+        out = tensor
+
+    src = get_group_rank(group, src)
+    group = _check_group(group)
+    return comm.broadcast(group, tensor, out, src, use_comm_stream, BroadcastAlgo.kNone)
+
+
+def reduce_scatter_tensor(
+    output: Tensor,
+    input: Tensor,
+    op=ReduceOp.SUM,
+    group: Optional[ProcessGroup] = None,
+    async_op=False,
+    use_comm_stream: bool = False,
+):
+    """
+    Example:
+        ixf_tensor_out = torch.zeros(2, dtype=torch.int64, device="cuda")
+        tensor_in = torch.arange(world_size * 2, dtype=torch.int64, device="cuda")
+        # tensor_in: tensor([0, 1, 2, 3, 4, 5, 6, 7], device='cuda:0')
+
+        ixfd.reduce_scatter_tensor(ixf_tensor_out, tensor_in, async_op=True)
+        print("rank {rank}:", ixf_tensor_out)
+
+        # output
+        rank 0:  tensor([0, 4], device='cuda:0')
+        rank 1:  tensor([ 8, 12], device='cuda:1')
+        rank 2:  tensor([16, 20], device='cuda:2')
+        rank 3:  tensor([24, 28], device='cuda:3')
+    """
+    if not async_op:
+        raise RuntimeError("Not support sync operation now.")
+
+    group = _check_group(group)
+    return comm.reduce_scatter(
+        group, input, output, op, use_comm_stream, ReduceScatterAlgo.kNone
+    )
+
+
+def all_reduce(
+    tensor: Tensor,
+    op=ReduceOp.SUM,
+    group: Optional[ProcessGroup] = None,
+    async_op=False,
+    out: Tensor = None,
+    algo: AllReduceAlgo = AllReduceAlgo.kNone,
+    use_comm_stream: bool = False,
+):
+    """
+    Args:
+        tensor: inpute tensor
+        op: ReduceOp: SUM, MIN or MAX
+        group: communicator group
+        async_op: ixformer support async mode
+        out: output tensor
+        algo: AllReduce Algo: Auto, Quant, QuantL1, QuantL2, NCCL, Ring, AllGatherSum, BroadcastSum
+        use_comm_stream: ixformer support set communication stream by ixformer.distributed.set_comm_group_stream,
+          if true, submit the kernels of communication to communication stream,
+          if false, use current stream by torch.cuda.current_stream
+    Returns: out
+
+    Example:
+        >>> # All tensors below are of torch.int64 type.
+        >>> # We have 2 process groups, 2 ranks.
+        >>> tensor = torch.arange(2, dtype=torch.int64) + 1 + 2 * rank
+        >>> tensor
+        tensor([1, 2]) # Rank 0
+        tensor([3, 4]) # Rank 1
+        >>> ixfd.all_reduce(tensor, op=ReduceOp.SUM, async_op=True)
+        >>> tensor
+        tensor([4, 6]) # Rank 0
+        tensor([4, 6]) # Rank 1
+    """
+    if not async_op:
+        raise RuntimeError("Not support sync operation now.")
+
+    group = _check_group(group)
+
+    if out is None:
+        out = tensor
+
+    comm.all_reduce(
+        group,
+        tensor,
+        out,
+        op,
+        use_comm_stream=use_comm_stream,
+        algo=algo,
+    )
+
+
+def all_gather_into_tensor(
+    output: Tensor,
+    input: Tensor,
+    group: Optional[ProcessGroup] = None,
+    async_op=False,
+    use_comm_stream: bool = False,
+):
+    """
+    Example:
+        tensor_in = torch.arange(2, dtype=torch.int64, device="cuda") + 1 + 2 * rank
+            rank 0:  tensor in:  tensor([1, 2], device='cuda:0')
+            rank 1:  tensor in:  tensor([3, 4], device='cuda:1')
+            rank 2:  tensor in:  tensor([5, 6], device='cuda:2')
+            rank 3:  tensor in:  tensor([7, 8], device='cuda:3')
+
+        ixf_tensor_out = torch.zeros(world_size * 2, dtype=torch.int64, device="cuda")
+        ixfd.all_gather_into_tensor(ixf_tensor_out, tensor_in, async_op=True)
+        print("rank {rank}:", ixf_tensor_out)
+
+        # output:
+        rank 0:  tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:0')
+        rank 1:  tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:1')
+        rank 2:  tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:2')
+        rank 3:  tensor([1, 2, 3, 4, 5, 6, 7, 8], device='cuda:3')
+    """
+    if not async_op:
+        raise RuntimeError("Not support sync operation now.")
+
+    group = _check_group(group)
+    return comm.all_gather(
+        group, input, output, use_comm_stream, algo=AllGatherAlgo.kNone
+    )
+
+
+def gather(
+    tensor,
+    gather_list=None,
+    dst=0,
+    group: Optional[ProcessGroup] = None,
+    async_op=False,
+    use_comm_stream: bool = False,
+):
+    """
+    Example:
+        >>> # We have 2 process groups, 2 ranks.
+        >>> tensor = torch.tensor(rank+1,dtype=torch.float32).cuda()
+        >>> tensor
+        tensor(1.) # Rank 0
+        tensor(2.) # Rank 1
+        >>> gather_list = [torch.zeros(1).cuda() for _ in range(rank)] if rank == dst else None
+        >>> gather_list
+        [tensor([0,]),tensor([1,])] # Rank 0
+        None                        # Rank 1
+        ixfd.gather(tensor,gather_list,0,async_op=True)
+        >>> gather_list
+        [tensor([1.]),tensor([2.])] # Rank 0
+        None                        # Rank 1
+    """
+    gather_list = gather_list if gather_list is not None else []
+    if not async_op:
+        raise RuntimeError("Not support sync operation now.")
+
+    dst = get_group_rank(group, dst)
+    group = _check_group(group)
+    return comm.gather(group, tensor, gather_list, dst, use_comm_stream)

ixformer_sdk/distributed/overlap_comm.py

@@ -0,0 +1,412 @@
+import abc
+import enum
+import os
+from contextlib import contextmanager, nullcontext
+from typing import List, Optional
+
+import torch.cuda
+from ixformer.core.dispatcher import Dispatcher
+
+from ixformer.core import config
+
+from . import _distributed as ixfd
+
+
+class SplitOverlapComm(Dispatcher):
+    def __init__(self, num_chunks, num_compute_streams=None, comm_group=None):
+        """
+        Args:
+            num_chunks: the number of chunks
+            num_compute_streams: the number of compute streams, default: 1
+            comm_group: communicator group
+        """
+
+        self._num_chunks = num_chunks
+        self._num_compute_streams = num_compute_streams or 1
+        self._comm_group = comm_group
+
+        self._compute_streams: List[torch.cuda.Stream] = self.create_compute_streams()
+        self._comm_stream: torch.cuda.Stream = torch.cuda.Stream(priority=-1)
+
+        self._start_compute_event: torch.cuda.Event = torch.cuda.Event()
+        self._stop_compute_event: torch.cuda.Event = torch.cuda.Event()
+
+        self._start_comm_event: torch.cuda.Event = torch.cuda.Event()
+        self._stop_comm_event: torch.cuda.Event = torch.cuda.Event()
+
+        # keep origin state
+        self._main_stream: Optional[torch.cuda.Stream] = None
+        self._origin_ixf_comm_stream = None
+        self._ixformer_streams = dict()
+
+    @classmethod
+    def dispatcher_key(
+        cls, num_chunks, num_compute_streams=None, comm_group=None, *args, **kwargs
+    ):
+        """
+        the key of SplitOverlapComm
+        Args:
+            num_chunks: the number of chunks
+            num_compute_streams: the number of compute streams, default: 1
+            comm_group: communicator group
+        Returns: unique key
+        """
+        # warn: keey same function parameters with init
+        return (cls.__name__, num_chunks, num_compute_streams, comm_group)
+
+    @classmethod
+    def enable(cls):
+        return config.IXFORMER_ENABLE_OVERLAP_COMM
+
+    @property
+    def num_chunks(self):
+        return self._num_chunks
+
+    @property
+    def num_compute_streams(self):
+        return self._num_compute_streams
+
+    @property
+    def comm_group(self):
+        return self._comm_group
+
+    def create_compute_streams(self):
+        streams = []
+        for _ in range(self.num_compute_streams):
+            streams.append(torch.cuda.Stream())
+        return streams
+
+    def start_overlap(self):
+        self._main_stream = torch.cuda.current_stream()
+
+        self._start_compute_event.record(torch.cuda.current_stream())
+        for compute_stream in self._compute_streams:
+            compute_stream.wait_event(self._start_compute_event)
+
+        self._origin_ixf_comm_stream = ixfd.get_comm_group_stream(self._comm_group)
+        ixfd.set_comm_group_stream(self._comm_stream.cuda_stream, self._comm_group)
+
+    def stop_overlap(self):
+        last_compute_stream_id = (
+            self.num_chunks + self.num_compute_streams - 1
+        ) % self.num_compute_streams
+        self._stop_compute_event.record(self._compute_streams[last_compute_stream_id])
+        self._stop_comm_event.record(self._comm_stream)
+        torch.cuda.current_stream().wait_event(self._stop_compute_event)
+        torch.cuda.current_stream().wait_event(self._stop_comm_event)
+
+        ixfd.set_comm_group_stream(self._origin_ixf_comm_stream, self._comm_group)
+
+    def start_comm(self, chunk_idx):
+        """
+        prepare communication stream and wait event.
+        Args:
+            chunk_idx: the index of chunk
+        """
+
+        self._start_comm_event.record(
+            self._compute_streams[chunk_idx % self.num_compute_streams]
+        )
+        self._comm_stream.wait_event(self._start_comm_event)
+
+    @contextmanager
+    def compute_stream_context(self, chunk_idx):
+        """
+        open python context and switch to compute stream in torch context
+        Args:
+            chunk_idx: the index of chunk
+        """
+
+        stream = self._compute_streams[chunk_idx % self.num_compute_streams]
+
+        # print("before stream:", torch.cuda.current_stream())
+        torch.cuda.set_stream(stream)
+
+        # print("after stream:", torch.cuda.current_stream(), ixformer.cuda.current_stream())
+        yield stream
+
+        torch.cuda.set_stream(self._main_stream)
+
+    @contextmanager
+    def stream_context(self, stream):
+        # print("before stream:", torch.cuda.current_stream())
+        torch.cuda.set_stream(stream)
+
+        # print("after stream:", torch.cuda.current_stream(), ixformer.cuda.current_stream())
+        yield stream
+
+        torch.cuda.set_stream(self._main_stream)
+
+    def forward(self, *args, **kwargs):
+        self.start_overlap()
+        out = self.compute(*args, **kwargs)
+        self.stop_overlap()
+        return out
+
+    @abc.abstractmethod
+    def compute(self, *args, **kwargs):
+        """
+        it is abstract method to execute compute and communication.
+        """
+        pass
+
+
+class GemmMethod(enum.IntEnum):
+    kCUINFER = 0
+    kCUBLAS = 1
+    kLIMITED_GEMM = 2
+
+
+class GemmWithLimitedBlock:
+    def __init__(self, limit_algo=0) -> None:
+        self.limit_algo = limit_algo
+        self.env_key = "PYTORCH_GEMM_BLOCK_LIMITATION"
+
+    def __enter__(self) -> None:
+        os.environ[self.env_key] = str(self.limit_algo)
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        del os.environ[self.env_key]
+
+
+class IxFormerLimitedGemmContext:
+    def __init__(self) -> None:
+        self.env_key = "IXFORMER_ENABLE_PERSISTENT_GEMM"
+
+    def __enter__(self) -> None:
+        os.environ[self.env_key] = "1"
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        os.environ[self.env_key] = "0"
+
+
+class GemmAllReduceSplitOverlapComm(SplitOverlapComm):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.gemm_method_env = config.IXFORMER_OVERLAP_GEMM_METHOD
+
+        if self.gemm_method_env is None:
+            if ixfd.get_world_size(self.comm_group) == 2:
+                self.gemm_method_env = 0
+            else:
+                self.gemm_method_env = 2
+
+        self.gemm_method = GemmMethod(int(self.gemm_method_env))
+        self.limited_gemm_ctx = GemmWithLimitedBlock()
+        self.ixf_limited_gemm_ctx = IxFormerLimitedGemmContext()
+        self.split_ratio = config.IXFORMER_OVERLAP_SPLIT_RATIO
+
+    @classmethod
+    def compute_row_parallel_dims(cls, input):
+        batch = 1
+        if input.ndim == 2:
+            seqlen = input.shape[0]
+        else:
+            batch = input.shape[0]
+            seqlen = input.shape[1]
+
+        parallel_dims = batch * seqlen
+        return parallel_dims
+
+    def compute(self, input, weight, bias=None, out=None, *args, **kwargs):
+        """
+        :param input: [Batch, SeqLen, Hidden]
+        :param weight: [OutChannel, InChannel]
+        :param bias: [OutChannel]
+        """
+
+        is_update_shape = input.ndim > 2
+        batch = 1
+        if input.ndim == 2:
+            seqlen = input.shape[0]
+        else:
+            batch = input.shape[0]
+            seqlen = input.shape[1]
+
+        parallel_dims = batch * seqlen
+
+        if is_update_shape:
+            input = input.reshape(parallel_dims, -1)
+
+        if out is None:
+            out_shape = [parallel_dims, weight.shape[0]]
+            out_dtype = kwargs["out_dtype"] if "out_dtype" in kwargs else input.dtype
+            out = torch.empty(out_shape, dtype=out_dtype, device=input.device)
+
+        if self.split_ratio is not None:
+            round_multiples = 256 if parallel_dims >= 256 else parallel_dims
+            first_chunk_size = (
+                round((parallel_dims * float(self.split_ratio)) / round_multiples)
+                * round_multiples
+            )
+            middle_chunk_size = (parallel_dims - first_chunk_size) // (
+                self.num_chunks - 1
+            )
+            middle_chunk_size = (middle_chunk_size // round_multiples) * round_multiples
+            last_chunk_size = (
+                parallel_dims
+                - first_chunk_size
+                - middle_chunk_size * (self.num_chunks - 2)
+            )
+
+            chunk_sizes = (
+                [first_chunk_size]
+                + [middle_chunk_size] * (self.num_chunks - 2)
+                + [last_chunk_size]
+            )
+            input_chunks = torch.split_with_sizes(input, chunk_sizes, dim=0)
+            out_chunks = torch.split_with_sizes(out, chunk_sizes, dim=0)
+
+            # print(first_chunk_size, middle_chunk_size, last_chunk_size, chunk_sizes)
+        else:
+            input_chunks = torch.chunk(input, self.num_chunks, dim=0)
+            out_chunks = torch.chunk(out, self.num_chunks, dim=0)
+
+        for chunk_idx in range(len(input_chunks)):
+            with self.compute_stream_context(chunk_idx):
+                chunk_out = self.gemm_dispatcher(
+                    chunk_idx,
+                    input_chunks[chunk_idx],
+                    weight,
+                    out_chunks[chunk_idx],
+                    *args,
+                    **kwargs,
+                )
+
+            self.start_comm(chunk_idx)
+
+            ixfd.all_reduce(
+                chunk_out, async_op=True, group=self.comm_group, use_comm_stream=True
+            )
+
+        if is_update_shape:
+            out = out.reshape(batch, seqlen, -1)
+
+        if bias is not None:
+            out = out + bias
+
+        return out
+
+    def gemm_dispatcher(
+        self,
+        chunk_idx,
+        chunk_input,
+        weight,
+        chunk_out=None,
+        user_gemm_method=None,
+        *args,
+        **kwargs,
+    ):
+        if user_gemm_method is not None and callable(user_gemm_method):
+            ctx = nullcontext() if chunk_idx == 0 else self.ixf_limited_gemm_ctx
+            with ctx:
+                return user_gemm_method(
+                    chunk_input, weight, out=chunk_out, *args, **kwargs
+                )
+
+        if user_gemm_method is None:
+            user_gemm_method = self.gemm_method
+
+        if user_gemm_method == GemmMethod.kCUINFER:
+            import ixformer.functions as ixff
+
+            return ixff.linear(chunk_input, weight, output=chunk_out)
+        elif user_gemm_method == GemmMethod.kCUBLAS:
+            return torch.matmul(chunk_input, weight.T, out=chunk_out)
+        elif user_gemm_method == GemmMethod.kLIMITED_GEMM:
+            ctx = self.limited_gemm_ctx
+            with ctx:
+                return torch.matmul(chunk_input, weight.T, out=chunk_out)
+        elif user_gemm_method == GemmMethod.kCUBLAS:
+            return torch.matmul(chunk_input, weight.T, out=chunk_out)
+        else:
+            raise RuntimeError(f"Invalid gemm method, got {self.gemm_method}.")
+
+    @classmethod
+    def native_forward(
+        cls,
+        input,
+        weight,
+        bias=None,
+        out=None,
+        group=None,
+        user_gemm_method=None,
+        *args,
+        **kwargs,
+    ):
+        if user_gemm_method is not None and callable(user_gemm_method):
+            gemm_out = user_gemm_method(
+                input, weight, bias=bias, out=out, *args, **kwargs
+            )
+            out = out if gemm_out is None else gemm_out
+        else:
+            import ixformer.functions as ixff
+
+            # warning: 下面的两种 gemm 可能存在精度不一致
+            # out = torch.matmul(input, weight.T, out=out)
+            out = ixff.linear(input=input, weight=weight, bias=bias, output=out)
+        ixfd.all_reduce(out, async_op=True, group=group)
+        return 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
+
+
+_DEFAULT_OVERLAP_GROUP = None
+_DEFAULT_OVERLAP_COMM_N2 = None
+_DEFAULT_OVERLAP_COMM_N4 = None
+_DEFAULT_OVERLAP_CHUNKS = config.IXFORMER_OVERLAP_CHUNKS
+
+
+def linear_allreduce_overlap(
+    input, weight, bias=None, out=None, group=None, num_chunks=None, *args, **kwargs
+):
+    num_chunks = num_chunks or _DEFAULT_OVERLAP_CHUNKS
+
+    # print("call overlap:", GemmAllReduceSplitOverlapComm.is_supported(input, num_chunks=num_chunks, comm_group=group), input.shape, weight.shape if torch.is_tensor(weight) else None, "WorldSize:", ixfd.get_group_world_size(group), ", NumChunks:", num_chunks)
+    if not GemmAllReduceSplitOverlapComm.is_supported(
+        input, num_chunks=num_chunks, comm_group=group
+    ):
+        return GemmAllReduceSplitOverlapComm.native_forward(
+            input, weight, bias=bias, out=out, group=group, *args, **kwargs
+        )
+
+    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 = GemmAllReduceSplitOverlapComm.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 = GemmAllReduceSplitOverlapComm.dispatcher(
+                num_chunks=num_chunks, comm_group=group
+            )
+        overlap_comm = _DEFAULT_OVERLAP_COMM_N4
+    else:
+        overlap_comm = GemmAllReduceSplitOverlapComm.dispatcher(
+            num_chunks=num_chunks, comm_group=group
+        )
+    return overlap_comm.forward(input, weight, bias=bias, out=out, *args, **kwargs)

ixformer_sdk/functions/__init__.py

@@ -0,0 +1 @@
+from ..inference.functions import *

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