chore: import upstream snapshot with attribution

deepspeed/inference/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+from .v2 import RaggedInferenceEngineConfig, DeepSpeedTPConfig
+from .v2.engine_v2 import InferenceEngineV2
+from .v2 import build_hf_engine, build_engine_from_ds_checkpoint

deepspeed/inference/config.py

@@ -0,0 +1,323 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed
+from pydantic import Field, field_validator
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
+from typing import Dict, Union, Optional
+from enum import Enum
+
+
+class DtypeEnum(Enum):
+    fp16 = (torch.float16, "torch.float16", "fp16", "float16", "half")
+    fp32 = (torch.float32, "torch.float32", "fp32", "float32", "float")
+    bf16 = (torch.bfloat16, "torch.bfloat16", "bf16", "bfloat16", "bfloat")
+    int8 = (torch.int8, "torch.int8", "int8")
+
+    @classmethod
+    def from_str(cls, value: str):
+        for dtype in cls:
+            if value in dtype.value:
+                return dtype
+        raise ValueError(f"'{value}' is not a valid DtypeEnum")
+
+
+class MoETypeEnum(str, Enum):
+    residual = "residual"
+    standard = "standard"
+
+
+class DeepSpeedTPConfig(DeepSpeedConfigModel):
+    """ Configure tensor parallelism settings """
+
+    enabled: bool = True
+    """ Turn tensor parallelism on/off. """
+
+    tp_size: int = 1
+    """ Number of devices to split the model across using tensor parallelism. """
+
+    tp_grain_size: int = 64
+    "Desired MLP/lm_head tp size granularity. DNN library favors tensor size in granularity of power of 2, we pick 64 as a default size."
+
+    mpu: object = None
+    """
+    A model parallelism unit object that implements
+    ``get_{model,data}_parallel_{rank,group,world_size}()``.
+    """
+
+    tp_group: object = None
+
+
+class DeepSpeedMoEConfig(DeepSpeedConfigModel):
+    """ Sets parameters for MoE """
+
+    enabled: bool = True
+    ep_size: int = 1
+    """
+    The expert-parallelism size which is used for partitioning the experts
+    across the GPUs in the expert-parallel group.
+    """
+
+    moe_experts: list = Field([1], alias="num_experts")
+    """ The global number of experts used in an MoE layer. """
+
+    type: MoETypeEnum = MoETypeEnum.standard
+    """
+    Specify the type of MoE layer. We have two types of MoE layer: 'Standard'
+    and 'Residual'.
+    """
+
+    ep_mp_group: object = None
+    ep_group: object = Field(None, alias="expert_group")
+
+
+class QuantTypeEnum(str, Enum):
+    asym = "asymmetric"
+    sym = "symmetric"
+
+
+class BaseQuantConfig(DeepSpeedConfigModel):
+    enabled: bool = True
+    num_bits: int = 8
+    q_type: QuantTypeEnum = QuantTypeEnum.sym
+    q_groups: int = 1
+
+
+class WeightQuantConfig(BaseQuantConfig):
+    enabled: bool = True
+    quantized_initialization: Dict = {}
+    post_init_quant: Dict = {}
+
+
+class ActivationQuantConfig(BaseQuantConfig):
+    enabled: bool = True
+
+
+class QKVQuantConfig(DeepSpeedConfigModel):
+    enabled: bool = True
+
+
+class QuantizationConfig(DeepSpeedConfigModel):
+    enabled: bool = True
+    activation: ActivationQuantConfig = ActivationQuantConfig()
+    weight: WeightQuantConfig = WeightQuantConfig()
+    qkv: QKVQuantConfig = QKVQuantConfig()
+
+
+# todo: brainstorm on how to do ckpt loading for DS inference
+class InferenceCheckpointConfig(DeepSpeedConfigModel):
+    checkpoint_dir: Optional[str] = None
+    save_mp_checkpoint_path: Optional[str] = None
+    base_dir: Optional[str] = None
+
+
+class DeepSpeedInferenceConfig(DeepSpeedConfigModel):
+    """ Sets parameters for DeepSpeed Inference Engine. """
+
+    replace_with_kernel_inject: bool = Field(False, alias="kernel_inject")
+    """
+    Set to true to inject inference kernels for models such as, Bert, GPT2,
+    GPT-Neo and GPT-J.  Otherwise, the injection_dict provides the names of two
+    linear layers as a tuple:
+    `(attention_output projection, transformer output projection)`
+    """
+
+    dtype: torch.dtype = torch.float16
+    """
+    Desired model data type, will convert model to this type.
+    Supported target types: `torch.half`, `torch.int8`, `torch.float`
+    """
+
+    tensor_parallel: DeepSpeedTPConfig = Field({}, alias="tp")
+    """
+    Configuration for tensor parallelism used to split the model across several
+    GPUs. Expects a dictionary containing values for :any:`DeepSpeedTPConfig`.
+    """
+
+    enable_cuda_graph: bool = False
+    """
+    Use this flag for capturing the CUDA-Graph of the inference ops, so that it
+    can run faster using the graph replay method.
+    """
+
+    use_triton: bool = False
+    """
+    Use this flag to use triton kernels for inference ops.
+    """
+
+    triton_autotune: bool = False
+    """
+    Use this flag to enable triton autotuning.
+    Turning it on is better for performance but increase the 1st runtime for
+    autotuning.
+    """
+
+    zero: DeepSpeedZeroConfig = {}
+    """
+    ZeRO configuration to use with the Inference Engine. Expects a dictionary
+    containing values for :any:`DeepSpeedZeroConfig`.
+    """
+
+    triangular_masking: bool = Field(True, alias="tm")
+    """
+    Controls the type of masking for attention scores in transformer layer.
+    Note that the masking is application specific.
+    """
+
+    moe: Union[bool, DeepSpeedMoEConfig] = {}
+    """
+    Specify if the type of Transformer is MoE. Expects a dictionary containing
+    values for :any:`DeepSpeedMoEConfig`.
+    """
+
+    keep_module_on_host: bool = False
+    """
+    When loading checkpoints to model parameters, they are moved to the device. In very large models
+    this might fill the device and cause OOM. Setting this flag to true, will keep checkpoints on
+    host and not move them directly to the device (giving an option to quantize checkpoint data before
+    moving it to the device for example).
+    Set only for models with injection policies and auto TP.
+    """
+
+    quant: QuantizationConfig = {}
+    """
+    NOTE: only works for int8 dtype.
+    Quantization settings used for quantizing your model using the MoQ.  The
+    setting can be one element or a tuple. If one value is passed in, we
+    consider it as the number of groups used in quantization. A tuple is passed
+    in if we want to mention that there is extra-grouping for the MLP part of a
+    Transformer layer (e.g. (True, 8) shows we quantize the model using 8
+    groups for all the network except the MLP part that we use 8 extra
+    grouping). Expects a dictionary containing values for
+    :any:`QuantizationConfig`.
+    """
+
+    #todo: refactor the following 3 into the new checkpoint_config
+    checkpoint: Optional[Union[str, Dict]] = None
+    """
+    Path to deepspeed compatible checkpoint or path to JSON with load policy.
+    """
+
+    base_dir: str = ""
+    """
+    This shows the root directory under which all the checkpoint files exists.
+    This can be passed through the json config too.
+    """
+
+    set_empty_params: bool = False
+    """
+    specifying whether the inference-module is created with empty or real Tensor
+    """
+
+    save_mp_checkpoint_path: Optional[str] = None
+    """
+    The path for which we want to save the loaded model with a checkpoint. This
+    feature is used for adjusting the parallelism degree to help alleviate the
+    model loading overhead. It does not save any new checkpoint if no path is
+    passed.
+    """
+
+    checkpoint_config: InferenceCheckpointConfig = Field({}, alias="ckpt_config")
+    """
+    TODO: Add docs. Expects a dictionary containing values for
+    :any:`InferenceCheckpointConfig`.
+    """
+
+    return_tuple: bool = True
+    """
+    Specify whether or not the transformer layers need to return a tuple or a
+    Tensor.
+    """
+
+    training_mp_size: int = 1
+    """
+    If loading a checkpoint this is the mp size that it was trained with, it
+    may be different than what the mp size that you want to use during
+    inference.
+    """
+
+    replace_method: str = Field(
+        "auto",
+        json_schema_extra={
+            "deprecated": True,
+            "deprecated_msg": "This parameter is no longer needed, please remove from your call to DeepSpeed-inference"
+        })
+
+    injection_policy: Optional[Dict] = Field(None, alias="injection_dict")
+    """
+    Dictionary mapping a client nn.Module to its corresponding injection
+    policy. e.g., `{BertLayer : deepspeed.inference.HFBertLayerPolicy}`
+    """
+
+    injection_policy_tuple: Optional[tuple] = None
+    """ TODO: Add docs """
+
+    config: Optional[Dict] = Field(None, alias="args")  # todo: really no need for this field if we can refactor
+
+    max_out_tokens: int = Field(1024, alias="max_tokens")
+    """
+    This argument shows the maximum number of tokens inference-engine can work
+    with, including the input and output tokens. Please consider increasing it
+    to the required token-length required for your use-case.
+    """
+
+    min_out_tokens: int = Field(1, alias="min_tokens")
+    """
+    This argument communicates to the runtime the minimum number of tokens you
+    expect you will need to generate. This will cause the runtime to error
+    if it unable to provide this and provide context on the memory pressure
+    rather than seg-faulting or providing corrupted output.
+    """
+
+    transposed_mode: bool = Field(False, alias="transposed_mode")
+
+    mp_size: int = Field(1, json_schema_extra={"deprecated": True, "new_param": "tensor_parallel.tp_size"})
+    """
+    Desired model parallel size, default is 1 meaning no model parallelism.
+    Deprecated, please use the ``tensor_parallel` config to control model
+    parallelism.
+    """
+    mpu: object = Field(None, json_schema_extra={"deprecated": True, "new_param": "tensor_parallel.mpu"})
+    ep_size: int = Field(1, json_schema_extra={"deprecated": True, "new_param": "moe.ep_size"})
+    ep_group: object = Field(None,
+                             alias="expert_group",
+                             json_schema_extra={
+                                 "deprecated": True,
+                                 "new_param": "moe.ep_group"
+                             })
+    ep_mp_group: object = Field(None,
+                                alias="expert_mp_group",
+                                json_schema_extra={
+                                    "deprecated": True,
+                                    "new_param": "moe.ep_mp_group"
+                                })
+    moe_experts: list = Field([1], json_schema_extra={"deprecated": True, "new_param": "moe.moe_experts"})
+    moe_type: MoETypeEnum = Field(MoETypeEnum.standard,
+                                  json_schema_extra={
+                                      "deprecated": True,
+                                      "new_param": "moe.type"
+                                  })
+
+    @field_validator("dtype", mode="before")
+    def validate_dtype(cls, field_value, values):
+        if isinstance(field_value, str):
+            return DtypeEnum.from_str(field_value).value[0]
+        if isinstance(field_value, torch.dtype):
+            return field_value
+        raise TypeError(f"Invalid type for dtype: {type(field_value)}")
+
+    @field_validator("moe")
+    def moe_backward_compat(cls, field_value, values):
+        if isinstance(field_value, bool):
+            return DeepSpeedMoEConfig(moe=field_value)
+        return field_value
+
+    @field_validator("use_triton")
+    def has_triton(cls, field_value, values):
+        if field_value and not deepspeed.HAS_TRITON:
+            raise ValueError('Triton needs to be installed to use deepspeed with triton kernels')
+        return field_value

deepspeed/inference/engine.py

@@ -0,0 +1,628 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import time
+import os
+import deepspeed
+from deepspeed import comm as dist
+from deepspeed.utils.logging import log_dist
+
+from torch.nn.modules import Module
+from packaging import version as pkg_version
+from deepspeed.runtime.checkpoint_engine.torch_checkpoint_engine import TorchCheckpointEngine
+from deepspeed.utils.timer import SynchronizedWallClockTimer
+from deepspeed.runtime.compiler import is_compile_supported
+from ..runtime.state_dict_factory import SDLoaderFactory
+from ..runtime.weight_quantizer import WeightQuantization
+from ..module_inject import replace_transformer_layer, generic_injection
+from ..comm.comm import init_distributed
+from ..pipe import PipelineModule
+from ..moe.utils import has_moe_layers
+from ..module_inject import LinearAllreduce, LinearLayer, Normalize, ReplaceWithTensorSlicing
+from deepspeed.accelerator import get_accelerator
+from ..module_inject.policy import TransformerPolicy
+from ..module_inject.auto_tp import AutoTP
+
+from ..module_inject.replace_policy import generic_policies
+from ..module_inject.auto_tp_model_utils import build_bloom_alibi_tensor, build_mpt_atten_bias_tensor, build_mpt_alibi_tensor, get_alibi_mask
+from ..ops.transformer.inference.ds_attention import DeepSpeedSelfAttention
+from ..model_implementations.transformers.ds_transformer import DeepSpeedTransformerInference
+
+DS_INFERENCE_ENABLED = False
+from torch import nn
+
+INFERENCE_MODEL_TIMER = "model-forward-inference"
+
+
+class InferenceEngine(Module):
+    inference_mp_group = None
+    inference_ep_group = None
+    expert_mp_group = None
+
+    def __init__(self, model, config):
+        """
+        Args:
+            model: torch.nn.Module
+            config: DeepSpeedInferenceConfig
+        """
+        global DS_INFERENCE_ENABLED
+        DS_INFERENCE_ENABLED = True
+
+        super().__init__()
+        if DeepSpeedTransformerInference.workspace is not None:
+            self.destroy()
+
+        self.module = model
+        self._config = config
+
+        self._get_model_config_generate(config)  # keep for weird backward compatibility
+
+        # patch model generate with ours if model uses it
+        if hasattr(self.module, "generate"):
+            self.generate = self._generate
+
+        if hasattr(self.module, "config"):
+            TransformerPolicy.hf_model_config = self.module.config
+
+        if config.dtype not in get_accelerator().supported_dtypes():
+            raise ValueError(
+                f"Data type {config.dtype} is not supported by {get_accelerator().device_name()} accelerator")
+
+        # todo: keep this self.injection_dict because we don't use to change config.injection_policy API
+        # todo: this will get changed when Molly's PR on auto injection dict is merged
+        self.injection_dict = config.injection_policy
+
+        # todo: refactor the mp_group and mp_size related in the next refactor
+        self.mp_group = config.tensor_parallel.tp_group
+        self.mpu = config.tensor_parallel.mpu
+
+        self.quantize_merge_count = 1
+        self.quantization_scales = None
+
+        # these are not needed in the config as we are creating them ourselves in the inference engine
+        self.ep_group = None  # config.moe.ep_group
+        self.expert_mp_group = None  # config.moe.ep_mp_group
+
+        self.cuda_graph_created = False
+        self.checkpoint_engine = TorchCheckpointEngine()
+        quantization_setting = None
+        self._init_quantization_setting(
+            quantization_setting)  # todo: update with the new quant config for weight quant
+        self.model_profile_enabled = False
+        self._model_times = []
+
+        if not self.injection_dict and config.replace_with_kernel_inject:
+            # This is a hack to remove the prepare_mask function on HF side for BLOOM architecture
+            self.remove_mask_prepare_for_bloom()
+
+        if self.injection_dict or not config.replace_with_kernel_inject:
+            # This is a hack to redefine the alibi func due to TP
+            if config.tensor_parallel.tp_size > 1:
+                self.build_alibi_tensor()
+                self.build_attn_bias()
+
+        if get_accelerator().device_name() == 'cuda' and config.enable_cuda_graph:
+            assert pkg_version.parse(torch.__version__) >= pkg_version.parse("1.10"), \
+                "If you want to use cuda graph, please upgrade torch to at least v1.10"
+
+        # convert model to intended dtype
+        if config.dtype:
+            self._convert_to_dtype(config)
+
+        if self.mpu:
+            config.tensor_parallel.tp_size = dist.get_world_size(group=self.mpu.get_model_parallel_group())
+            self.mp_group = self.mpu.get_model_parallel_group()
+        elif config.tensor_parallel.tp_size > 1:
+            self._create_model_parallel_group(config)
+            config.tensor_parallel.tp_group = self.mp_group
+
+        if isinstance(self.module, torch.nn.Module):
+            moe, _ = has_moe_layers(self.module)
+        else:
+            moe = False
+
+        if moe and dist.get_world_size() > 1:
+            self._create_ep_parallel_group(config.moe.ep_size)
+
+        # We only support three modes: 1) user specified policy for tensor-parallelism, 2) kernel injection (replace_with_kernel_inject), and 3) automatic tensor parallelism if tp_size > 1.
+        if self.injection_dict:
+            # 1. User specified Tensor Parallelism
+            assert not config.replace_with_kernel_inject, "Cannot use both user specified injection policy and kernel injection"
+            for client_module, injection_policy in self.injection_dict.items():
+
+                assert issubclass(client_module,
+                                  torch.nn.Module), f"{client_module} is not a subclass of torch.nn.Module"
+
+                # construct the tuple and pass that instead of a string or dict.
+                if isinstance(injection_policy, str):
+                    config.injection_policy_tuple = (injection_policy, )
+                else:
+                    config.injection_policy_tuple = injection_policy
+
+                layer_names = [name for name, _ in self.module.named_modules()]
+                for policy in config.injection_policy_tuple:
+                    if not any(name.endswith(policy) for name in layer_names):
+                        raise ValueError(f"Injection policy layer'{policy}' not valid.")
+
+                self._apply_injection_policy(config, client_module)
+        else:
+            if config.replace_with_kernel_inject:
+                # 2. DeepSpeed Kernel Injection
+                self._apply_injection_policy(config)
+            elif config.tensor_parallel.tp_size > 1:
+                # 3. Automatic Tensor Parallelism
+                parser_dict = AutoTP.tp_parser(model)
+                print("AutoTP: ", parser_dict)
+                for client_module, injection_policy in parser_dict:
+                    if isinstance(injection_policy, str):
+                        config.injection_policy_tuple = (injection_policy, )
+                    else:
+                        config.injection_policy_tuple = injection_policy
+                    self._apply_injection_policy(config, client_module)
+
+        device = get_accelerator().current_device_name()
+        # NOTE: This check assumes a Hugging Face hierarchy for the device type i.e. module.device.type
+        is_meta_device = hasattr(self.module, "device") and self.module.device.type == 'meta'
+        if is_meta_device:
+            self.module.to_empty(device=device)
+        elif not config.keep_module_on_host:
+            self.module.to(device)
+
+        if config.tensor_parallel.tp_size > 1:
+            _rng_state = get_accelerator().get_rng_state().to(get_accelerator().current_device_name())
+            dist.broadcast(_rng_state, 0)
+            get_accelerator().set_rng_state(_rng_state.cpu())
+
+        if config.tensor_parallel.tp_size > 1:
+            assert not config.enable_cuda_graph, "Cuda graph is not supported for model parallelism"
+
+        # Check if local CUDA graphs can be created in replacement modules
+        self.local_cuda_graph = self._local_cuda_graph_used(self.module)
+        self._is_compiled = False
+
+    def destroy(self):
+        DeepSpeedTransformerInference.layer_id = 0
+        DeepSpeedSelfAttention.num_layers = 0
+        if DeepSpeedTransformerInference.workspace.is_allocated():
+            DeepSpeedTransformerInference.workspace.release_workspace()
+        DeepSpeedTransformerInference.workspace = None
+
+    def profile_model_time(self, use_cuda_events=True):
+        if not self.model_profile_enabled and not self._config.enable_cuda_graph:
+            self.module.register_forward_pre_hook(self._pre_forward_hook)
+            self.module.register_forward_hook(self._post_forward_hook)
+        self.model_profile_enabled = True
+        self.use_cuda_events = use_cuda_events
+        if self.use_cuda_events:
+            self.timers = SynchronizedWallClockTimer()
+
+    # todo: remove this once all the config dicts are centralized from top level pydantic config
+    def _get_model_config_generate(self, config):
+        # this is being passed to replace_transformer_layer(config=self.user_model_config_dict)
+        self.config = getattr(self.module, 'config', None) if config.config is None else config.config
+
+    def remove_mask_prepare_for_bloom(self):
+        if hasattr(self.module, 'transformer'):
+            if hasattr(self.module.transformer, '_prepare_attn_mask'):
+                self.module.transformer._prepare_attn_mask = lambda attention_mask, *args, **kwargs: attention_mask
+
+    def build_alibi_tensor(self):
+        if hasattr(self.module, 'transformer'):
+            if hasattr(self.module.transformer, 'build_alibi_tensor'):
+                self.module.transformer.build_alibi_tensor = build_bloom_alibi_tensor
+            if hasattr(self.module.transformer, 'build_mpt_alibi_tensor'):
+                self.module.transformer.build_mpt_alibi_tensor_orig = self.module.transformer.build_mpt_alibi_tensor
+                self.module.transformer.__class__.build_mpt_alibi_tensor = build_mpt_alibi_tensor
+        if hasattr(self.module, 'model'):
+            if hasattr(self.module.model, 'get_alibi_mask'):
+                self.module.model.get_alibi_mask_orig = self.module.model.get_alibi_mask
+                self.module.model.__class__.get_alibi_mask = get_alibi_mask
+
+    def build_attn_bias(self):
+        if hasattr(self.module, 'transformer'):
+            if hasattr(self.module.transformer, '_attn_bias'):
+                self.module.transformer._attn_bias_orig = self.module.transformer._attn_bias
+                self.module.transformer.__class__._attn_bias = build_mpt_atten_bias_tensor
+
+    def _pre_forward_hook(self, module, *inputs, **kwargs):
+        if self.use_cuda_events:
+            self.timers(INFERENCE_MODEL_TIMER).start()
+        else:
+            get_accelerator().synchronize()
+            self._start = time.time()
+
+    def _post_forward_hook(self, module, input, output):
+        if self.use_cuda_events:
+            self.timers(INFERENCE_MODEL_TIMER).stop()
+            elapsed_time = self.timers(INFERENCE_MODEL_TIMER).elapsed(reset=True)
+        else:
+            get_accelerator().synchronize()
+            self._end = time.time()
+            elapsed_time = (self._end - self._start) * 1e3  # convert seconds to ms
+        self._model_times.append(elapsed_time)
+
+    def _create_model_parallel_group(self, config):
+        # Call the init process
+        if InferenceEngine.inference_mp_group is None:
+            init_distributed()
+            local_rank = int(os.getenv('LOCAL_RANK', '0'))
+            get_accelerator().set_device(local_rank)
+
+            ranks = [i for i in range(config.tensor_parallel.tp_size)]
+            self.mp_group = dist.new_group(ranks)
+            InferenceEngine.inference_mp_group = self.mp_group
+        else:
+            self.mp_group = InferenceEngine.inference_mp_group
+
+    def _create_ep_parallel_group(self, moe_experts):
+        # Call the init process
+        self.ep_group = {}
+        self.expert_mp_group = {}
+        moe_experts = moe_experts if type(moe_experts) is list else [moe_experts]
+        for e in moe_experts:
+            self.ep_group.update({e: None})
+            self.expert_mp_group.update({e: None})
+        for moe_ep_size in self.ep_group.keys():
+            num_ep_groups = dist.get_world_size() // moe_ep_size
+            if num_ep_groups == 0:
+                raise ValueError(f"Invalid ep_size={moe_ep_size} for world_size={dist.get_world_size()}")
+            for i in range(num_ep_groups):
+                ep_cnt = i * moe_ep_size
+                size = dist.get_world_size() if moe_ep_size > dist.get_world_size() else moe_ep_size
+                ranks = list(range(ep_cnt, ep_cnt + size))
+                _ep_group = dist.new_group(ranks)
+                if dist.get_rank() in ranks:
+                    self.ep_group.update({moe_ep_size: _ep_group})
+
+            if dist.get_world_size() > moe_ep_size:
+                num_expert_mp_groups = dist.get_world_size() // num_ep_groups
+                expert_mp_size = dist.get_world_size() // moe_ep_size
+                for i in range(num_expert_mp_groups):
+                    expert_mp_comm_ranks = [i + nr * moe_ep_size for nr in range(expert_mp_size)]
+                    _expert_mp_group = dist.new_group(expert_mp_comm_ranks)
+                    if dist.get_rank() in expert_mp_comm_ranks:
+                        self.expert_mp_group.update({moe_ep_size: _expert_mp_group})
+
+    def _init_quantization_setting(self, quantization_setting):
+        self.quantize_bits = 8
+        self.mlp_extra_grouping = False
+        self.quantize_groups = 1
+        if type(quantization_setting) is tuple:
+            self.mlp_extra_grouping, \
+            self.quantize_groups = quantization_setting
+        elif quantization_setting is not None:
+            self.quantize_groups = quantization_setting
+        log_dist(
+            f"quantize_bits = {self.quantize_bits} "
+            f"mlp_extra_grouping = {self.mlp_extra_grouping}, "
+            f"quantize_groups = {self.quantize_groups}", [0])
+
+    def load_model_with_checkpoint(self, r_module):
+        self.mp_replace = ReplaceWithTensorSlicing(
+            mp_group=self.mp_group, mp_size=self._config.tensor_parallel.tp_size)  #, out_dim=0, in_dim=1)
+        error_msgs = []
+
+        def load(module, state_dict, prefix):
+            args = (state_dict, prefix, {}, True, [], [], error_msgs)
+            if hasattr(module, 'weight'):
+                if module.weight.data.is_meta:
+                    # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
+                    module.weight = torch.nn.parameter.Parameter(data=torch.empty_like(module.weight.data,
+                                                                                       device="cpu"),
+                                                                 requires_grad=module.weight.data.requires_grad)
+                if 'query_key_value' in prefix:
+                    module.weight = self.mp_replace.strided_copy(module.weight.data,
+                                                                 state_dict[prefix + 'weight'],
+                                                                 num_splits=3)
+                else:
+                    module.weight = self.mp_replace.copy(module.weight.data, state_dict[prefix + 'weight'])
+            else:
+                if module.norm.weight.data.is_meta:
+                    # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
+                    module.norm.weight = torch.nn.parameter.Parameter(
+                        data=torch.empty_like(module.norm.weight.data, device="cpu"),
+                        requires_grad=module.norm.weight.data.requires_grad)
+                module.norm.weight = self.mp_replace.copy(module.norm.weight.data, state_dict[prefix + 'weight'])
+            if prefix + 'bias' in self.key_list:
+                if hasattr(module, 'norm'):
+                    if module.norm.bias.data.is_meta:
+                        # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
+                        module.norm.bias = torch.nn.parameter.Parameter(
+                            data=torch.empty_like(module.norm.bias.data, device="cpu"),
+                            requires_grad=module.norm.bias.data.requires_grad)
+                    module.norm.bias = self.mp_replace.copy(module.norm.bias, state_dict[prefix + 'bias'])
+                else:
+                    if module.bias.data.is_meta:
+                        # meta tensor cannot be casted or copied to, so we need to replace it with a normal tensor here
+                        module.bias = torch.nn.parameter.Parameter(data=torch.empty_like(module.bias.data,
+                                                                                         device="cpu"),
+                                                                   requires_grad=module.bias.data.requires_grad)
+                    data = state_dict[prefix + 'bias']
+                    data = data.to(get_accelerator().current_device_name())
+                    module.bias = self.mp_replace.copy(module.bias, data)
+
+        layer_policies = {
+            nn.Linear: load,
+            nn.Embedding: load,
+            nn.LayerNorm: load,
+            LinearLayer: load,
+            LinearAllreduce: load
+        }
+
+        def load_module_recursive(module, prefix='', level=0):
+            for name, child in module.named_children():
+                if child.__class__ in layer_policies:
+                    checking_key = prefix + name + '.'
+                    if not any(checking_key in item for item in self.key_list):
+                        continue
+                    if len(list(child.parameters())) > 0 and list(child.parameters())[0].numel() == 0:
+                        if len(child.weight.ds_shape) == 1:
+                            child = Normalize(dim=child.weight.ds_shape[-1], dtype=child.weight.dtype, eps=child.eps)
+                            setattr(module, name, child)
+                    load(child, self.sd, prefix + name + '.')
+                else:
+                    load_module_recursive(child, prefix if level == 0 else prefix + name + '.', level + 1)
+
+        load_module_recursive(r_module)
+
+        embedding_weight = None
+
+        for n, p in r_module.named_parameters():
+            if "word_embeddings." in n or "embed_tokens." in n or "wte." in n:
+                embedding_weight = p
+        if embedding_weight is not None and hasattr(r_module, "lm_head") and hasattr(
+                r_module.lm_head, "weight") and r_module.lm_head.weight.is_meta:
+            r_module.lm_head.weight = embedding_weight
+
+    def _apply_injection_policy(self, config, client_module=None):
+        # client_module is only passed when using the injection_dict method.
+        checkpoint_dir = config.checkpoint
+        checkpoint = SDLoaderFactory.get_sd_loader_json(checkpoint_dir,
+                                                        self.checkpoint_engine) if checkpoint_dir is not None else None
+
+        generic_injection(self.module, dtype=config.dtype, enable_cuda_graph=config.enable_cuda_graph)
+
+        if isinstance(self.module, torch.nn.Module):
+            # config is our DeepSpeedInferenceConfig and self.config is the HF model config
+            replace_transformer_layer(client_module, self.module, checkpoint, config, self.config)
+
+    def _get_all_ckpt_names(self, checkpoints_path, tag):
+        ckpt_file_pattern = self._get_ckpt_name(checkpoints_path, tag, mp_placeholder="*")
+        import glob
+
+        ckpt_files = glob.glob(ckpt_file_pattern)
+        ckpt_files.sort()
+        return ckpt_files
+
+    def _get_ckpt_name(self, checkpoints_path, tag, mp_placeholder=None):
+        if mp_placeholder is not None:
+            mp_rank_str = mp_placeholder
+        else:
+            mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
+            mp_rank_str = "{:02d}".format(mp_rank)
+
+        ckpt_name = os.path.join(
+            checkpoints_path,
+            "mp_rank_" + mp_rank_str + "_model_states.pt",
+        )
+        return ckpt_name
+
+    def _load_checkpoint(self, load_dir, load_module_strict=True, tag=None):
+        is_pipe_parallel = isinstance(self.module, PipelineModule)
+        if is_pipe_parallel:
+            raise RuntimeError('pipeline parallelism is currently not supported in inference.')
+        if not isinstance(load_dir, dict) and os.path.isdir(load_dir):
+            if tag is None:
+                latest_path = os.path.join(load_dir, "latest")
+                if os.path.isfile(latest_path):
+                    with open(latest_path, "r") as fd:
+                        tag = fd.read().strip()
+
+            ckpt_list = self._get_all_ckpt_names(load_dir, tag)
+            sd_loader = SDLoaderFactory.get_sd_loader(ckpt_list, self.checkpoint_engine)
+        else:
+            sd_loader = SDLoaderFactory.get_sd_loader_json(load_dir, self.checkpoint_engine)
+
+        checkpoint = sd_loader['checkpoints']
+
+        if type(checkpoint) is list:
+            self.sd = torch.load(checkpoint[0], map_location='cpu', weights_only=False)
+            self.key_list = list(self.sd.keys())
+
+            self.load_model_with_checkpoint(self.module)
+
+            for i in range(1, len(checkpoint)):
+                if not dist.is_initialized() or dist.get_rank() == 0:
+                    print(f"loading checkpoint ({i})")
+                self.sd = torch.load(checkpoint[i], map_location=get_accelerator().device_name(), weights_only=False)
+                self.key_list = list(self.sd.keys())
+                self.load_model_with_checkpoint(self.module)
+        else:
+            mp_rank = 0 if self.mpu is None else self.mpu.get_model_parallel_rank()
+
+            load_path, checkpoint, quantize_config = sd_loader.load(self._config.tensor_parallel.tp_size,
+                                                                    mp_rank,
+                                                                    is_pipe_parallel=is_pipe_parallel,
+                                                                    quantize=(self._config.dtype is torch.int8),
+                                                                    quantize_groups=self.quantize_groups,
+                                                                    mlp_extra_grouping=self.mlp_extra_grouping)
+
+            self.quantization_scales, self.quantize_merge_count = quantize_config
+
+            moe, _ = has_moe_layers(self.module)
+            if moe:
+                from deepspeed.runtime.engine import DeepSpeedEngine
+                old_moe_load = False
+                if not isinstance(checkpoint['num_experts'], list):
+                    old_moe_load = True
+                DeepSpeedEngine.load_moe_state_dict(load_dir,
+                                                    tag,
+                                                    state_dict=checkpoint[self._choose_module_key(checkpoint)],
+                                                    old_moe_load=old_moe_load,
+                                                    model=self.module,
+                                                    mpu=self.mpu,
+                                                    checkpoint_engine=self.checkpoint_engine,
+                                                    autoep_layers=None)
+
+            self.module.load_state_dict(state_dict=checkpoint[self._choose_module_key(checkpoint)],
+                                        strict=load_module_strict)
+
+    def _choose_module_key(self, sd):
+        assert not ('module' in sd
+                    and 'model' in sd), "checkpoint has both 'model' and 'module' keys, not sure how to proceed"
+        assert 'module' in sd or 'model' in sd, "checkpoint contains neither 'model' or 'module' keys, not sure how to proceed"
+        if 'module' in sd:
+            return 'module'
+        elif 'model' in sd:
+            return 'model'
+
+    def _convert_to_dtype(self, config):
+        if not isinstance(self.module, torch.nn.Module):
+            return
+
+        if False:  #config.dtype is torch.int8 and self.quantization_scales is None:
+            quantizer = WeightQuantization(mlp_extra_grouping=self.mlp_extra_grouping)
+            model, self.quantization_scales = quantizer.model_quantize(self.module, self.injection_dict,
+                                                                       self.quantize_bits, self.quantize_groups)
+        elif config.dtype == torch.half:
+            self.module.half()
+        elif config.dtype == torch.bfloat16:
+            self.module.bfloat16()
+        elif config.dtype == torch.float:
+            self.module.float()
+
+    def _create_cuda_graph(self, *inputs, **kwargs):
+        # warmup to create the workspace and cublas handle
+        cuda_stream = get_accelerator().Stream()
+        cuda_stream.wait_stream(get_accelerator().current_stream())
+        with get_accelerator().stream(cuda_stream):
+            for i in range(3):
+                ret = self.module(*inputs, **kwargs)
+        get_accelerator().current_stream().wait_stream(cuda_stream)
+
+        # create cuda_graph and assign static_inputs and static_outputs
+        self._cuda_graphs = get_accelerator().create_graph()
+        self.static_inputs = inputs
+        self.static_kwargs = kwargs
+
+        with get_accelerator().capture_to_graph(self._cuda_graphs):
+            self.static_output = self.module(*self.static_inputs, **self.static_kwargs)
+
+        self.cuda_graph_created = True
+
+    def _graph_replay(self, *inputs, **kwargs):
+        for i in range(len(inputs)):
+            if torch.is_tensor(inputs[i]):
+                self.static_inputs[i].copy_(inputs[i])
+        for k in kwargs:
+            if torch.is_tensor(kwargs[k]):
+                self.static_kwargs[k].copy_(kwargs[k])
+        get_accelerator().replay_graph(self._cuda_graphs)
+        return self.static_output
+
+    def model_times(self):
+        assert self.model_profile_enabled, "model profiling is not enabled"
+        model_times = self._model_times
+        if self._config.enable_cuda_graph and len(self._model_times) == 0:
+            raise ValueError("Model times are empty and cuda graph is enabled. If "
+                             "this is a GPT-style model this combo is not supported. If this is a "
+                             "BERT-style model this is a bug, please report it. "
+                             f"Model type is: {type(self.module)}")
+        self._model_times = []
+        return model_times
+
+    def _module_match(self, module):
+        for policy in generic_policies:
+            policy = policy()
+            if policy.match_replaced(module):
+                return True
+        return False
+
+    def _local_cuda_graph_used(self, module):
+        if isinstance(module, torch.nn.Module):
+            return False
+        else:
+            sub_module_cuda_graph = False
+            for name in module.__dict__.keys():
+                sub_module = getattr(module, name)
+
+                if self._module_match(sub_module) and hasattr(sub_module, "enable_cuda_graph"):
+                    sub_module_cuda_graph = True
+
+            return sub_module_cuda_graph
+
+    def forward(self, *inputs, **kwargs):
+        """Execute forward propagation
+
+        Arguments:
+            *inputs: Variable length input list
+            **kwargs: variable length keyword arguments
+        """
+        start = None
+        if self.model_profile_enabled and get_accelerator().device_name() == 'cuda' and self._config.enable_cuda_graph:
+            get_accelerator().synchronize()
+            start = time.time()
+
+        if get_accelerator().device_name() == 'cuda' and self._config.enable_cuda_graph and not self.local_cuda_graph:
+            if self.cuda_graph_created:
+                outputs = self._graph_replay(*inputs, **kwargs)
+            else:
+                self._create_cuda_graph(*inputs, **kwargs)
+                outputs = self._graph_replay(*inputs, **kwargs)
+
+        else:
+            outputs = self.module(*inputs, **kwargs)
+
+        if self.model_profile_enabled and self._config.enable_cuda_graph:
+            get_accelerator().synchronize()
+            duration = (time.time() - start) * 1e3  # convert seconds to ms
+            self._model_times.append(duration)
+
+        return outputs
+
+    def _generate(self, *inputs, **kwargs):
+        # Reset KV-cache at the beginning of generate
+        if hasattr(self.module, 'reset_cache'):
+            self.module.reset_cache()
+        num_beams = 1
+        if "generation_config" in kwargs:
+            gen_config = kwargs["generation_config"]
+            num_beams = getattr(gen_config, "num_beams", 1)
+        if "num_beams" in kwargs:
+            num_beams = kwargs["num_beams"]
+
+        if num_beams > 1:
+            raise NotImplementedError("DeepSpeed does not support `num_beams` > 1, if this is important to you please "
+                                      "add your request to: https://github.com/deepspeedai/DeepSpeed/issues/2506")
+
+        if ("input_ids" in kwargs) and (kwargs["input_ids"].dim() == 2):
+            for input_tensor in kwargs["input_ids"]:
+                tensor_length = input_tensor.shape[-1]
+                if tensor_length > self._config.max_out_tokens:
+                    raise RuntimeError(
+                        f"Input with size {tensor_length} exceeds maximum length of {self._config.max_out_tokens}. Please increase max_tokens in the DeepSpeed Inference Config."
+                    )
+
+        return self.module.generate(*inputs, **kwargs)
+
+    def compile(self, backend=get_accelerator().get_compile_backend(), compile_kwargs={}) -> None:
+        """
+        Compile the module using the specified backend and kwargs.
+        """
+        if not is_compile_supported():
+            raise RuntimeError("compile is not supported in your version of PyTorch.")
+
+        if self._is_compiled:
+            return
+
+        # Avoid graph breaks
+        deepspeed.utils.nvtx.enable_nvtx = False
+        self.module.compile(backend=backend, **compile_kwargs)
+        self._is_compiled = True
+
+    @property
+    def is_compiled(self) -> bool:
+        return self._is_compiled

deepspeed/inference/quantization/__init__.py

@@ -0,0 +1,4 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team

deepspeed/inference/quantization/layers.py

@@ -0,0 +1,114 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from torch import nn
+from torch import Tensor
+from torch.nn import functional as F
+from .utils import Quantizer, DeQuantizer, concat_to_compat_param
+from typing import Tuple, Callable, Dict
+from deepspeed.runtime.zero import register_external_parameter
+
+quantized_weight_registry = {}
+is_zero3_enabled = False
+
+
+# deal with weight sharing
+def get_quantized_weight_wrapper(model, pre_quant_weight: nn.Parameter, quantize_weight_fn: Callable) -> nn.Parameter:
+    if id(pre_quant_weight) in quantized_weight_registry:
+        compat_tensor = quantized_weight_registry[id(pre_quant_weight)]
+        if is_zero3_enabled:
+            register_external_parameter(model, compat_tensor)
+
+        return quantized_weight_registry[id(pre_quant_weight)]
+    else:
+        quantized_weights, quant_scale, quant_min = quantize_weight_fn()
+        quantized_weight_registry[id(pre_quant_weight)] = concat_to_compat_param(quantized_weights, quant_scale,
+                                                                                 quant_min)
+        return quantized_weight_registry[id(pre_quant_weight)]
+
+
+def get_quantize_weight_fn(quantizer: Quantizer, pre_quant_weight: nn.Parameter) -> Callable:
+
+    def func() -> Tuple[nn.Parameter, Tensor, Tensor]:
+        quantized_weights, quant_scale, quant_min = quantizer.quantize(pre_quant_weight.data)
+        # A temporary hack as zero Zero3 assume all model weights has the same type. in all_gather_coalesced.get_only_unique_item
+        quantized_weights = quantized_weights.view(pre_quant_weight.dtype)
+        quant_scale = quant_scale.type(pre_quant_weight.dtype)
+        quant_min = quant_min.type(pre_quant_weight.dtype)
+        return quantized_weights, quant_scale, quant_min
+
+    return func
+
+
+class QuantizedLinear(nn.Linear):
+
+    def __init__(self, config: Dict, pre_quant_layer: nn.Linear) -> None:
+        super(QuantizedLinear, self).__init__(in_features=pre_quant_layer.in_features,
+                                              out_features=pre_quant_layer.out_features,
+                                              bias=pre_quant_layer.bias is not None,
+                                              device=pre_quant_layer.weight.device,
+                                              dtype=pre_quant_layer.weight.dtype)
+        self.config = config
+
+        self.quantizer = Quantizer(config=config)
+        self.bias = pre_quant_layer.bias
+        self.weight = get_quantized_weight_wrapper(self, pre_quant_layer.weight,
+                                                   get_quantize_weight_fn(self.quantizer, pre_quant_layer.weight))
+
+        self.weight.dequantizer = DeQuantizer(config, pre_quant_layer.weight.dtype)
+
+    def forward(self, input: Tensor) -> Tensor:
+        quantized_weight, quant_scale, quant_min = self.weight.deconcat(self.weight)
+        temp_dequantized_weight = self.weight.dequantizer.dequantize(quantized_weight.view(torch.uint8), quant_scale,
+                                                                     quant_min)
+
+        # !!! Do not use torch.functional.linear(input, temp_dequantized_weight, self.bias) here as in zero3 torch.functional.linear is
+        # replaced by LinearFunctionForZeroStage3. Which assume weight is non-temporary.
+        # If weight is temp buffer there will be memory leak.
+        return torch._C._nn.linear(input, temp_dequantized_weight, self.bias)
+
+
+class QuantizedEmbedding(nn.Embedding):
+
+    def __init__(self, config: Dict, pre_quant_layer: nn.Embedding) -> None:
+        super(QuantizedEmbedding, self).__init__(num_embeddings=pre_quant_layer.num_embeddings,
+                                                 embedding_dim=pre_quant_layer.embedding_dim,
+                                                 padding_idx=pre_quant_layer.padding_idx,
+                                                 max_norm=pre_quant_layer.max_norm,
+                                                 norm_type=pre_quant_layer.norm_type,
+                                                 scale_grad_by_freq=pre_quant_layer.scale_grad_by_freq,
+                                                 sparse=pre_quant_layer.sparse,
+                                                 _weight=pre_quant_layer.weight,
+                                                 device=pre_quant_layer.weight.device,
+                                                 dtype=pre_quant_layer.weight.dtype)
+
+        assert pre_quant_layer.max_norm is None, 'Not supported'
+        assert pre_quant_layer.norm_type == 2, 'Not supported'
+        assert pre_quant_layer.scale_grad_by_freq == False, 'Not supported'
+        assert pre_quant_layer.sparse == False, 'Not supported'
+
+        self.config = config
+        quantizer = Quantizer(config=config)
+
+        self.weight = get_quantized_weight_wrapper(self, pre_quant_layer.weight,
+                                                   get_quantize_weight_fn(quantizer, pre_quant_layer.weight))
+
+        self.weight.dequantizer = DeQuantizer(config, pre_quant_layer.weight.dtype)
+
+    def forward(self, input: Tensor) -> Tensor:
+        quantized_weight, quant_scale, quant_min = self.weight.deconcat(self.weight)
+        temp_dequantized_weight = self.weight.dequantizer.dequantize(quantized_weight.view(torch.uint8), quant_scale,
+                                                                     quant_min)
+
+        return F.embedding(input, temp_dequantized_weight, self.padding_idx, self.max_norm, self.norm_type,
+                           self.scale_grad_by_freq, self.sparse)
+
+
+QUANTIZATION_LAYER_MAPPINGS = {
+    nn.Linear: QuantizedLinear,
+    nn.Embedding: QuantizedEmbedding,
+}

deepspeed/inference/quantization/quantization.py

@@ -0,0 +1,111 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from torch import nn
+from typing import Dict
+import gc
+from deepspeed.inference.quantization import layers
+from .layers import QUANTIZATION_LAYER_MAPPINGS
+from .utils import get_AsyncPartitionedParameterSwapper, recursive_setattr
+from deepspeed.utils.logging import logger
+from collections import deque
+from transformers.utils.generic import ContextManagers
+from .quantization_context import QuantizationContext
+import contextlib
+
+
+def _init_group_wise_weight_quantization(model: nn.Module, ds_config: Dict) -> nn.Module:
+    """[Experimental] Apply group-wise weight quantization to model. Replace layers module according to config_list
+
+    Args:
+        model (nn.Module): A nn.Module
+        ds_config (Dict, optional): The ds_config dictionary. use None for non-deepspeed managed model.
+
+    Returns:
+        nn.Module: Quantized nn.Module
+    """
+
+    # global quantized_weight_registry
+
+    matched_module_list_by_key = {}
+    matched_module_count = 0
+
+    assert 'weight_quantization' in ds_config, 'Please provide quantization config in ds_config'
+    quantization_config = ds_config['weight_quantization']['post_init_quant']
+
+    # Return nvme swapper if exists, else return None.
+    # For nvme offloading we must use the same swapper here as model initialized.
+    nvme_swapper = get_AsyncPartitionedParameterSwapper(model)
+    is_zero3_enabled = 'zero_optimization' in ds_config and \
+            'stage' in ds_config['zero_optimization'] and \
+            ds_config['zero_optimization']['stage'] == 3
+    is_offloading_enabled = 'zero_optimization' in ds_config and \
+                            'offload_param' in ds_config['zero_optimization']
+
+    layers.is_zero3_enabled = is_zero3_enabled
+
+    context_mgr = ContextManagers([QuantizationContext(config_dict_or_path=ds_config, param_swapper=nvme_swapper)]) \
+                    if is_zero3_enabled else contextlib.suppress()
+    with context_mgr:
+        module_list = list(
+            filter(lambda named_module: type(named_module[1]) in QUANTIZATION_LAYER_MAPPINGS, model.named_modules()))
+
+        # Quantize small weight first then large.
+        if not is_offloading_enabled:
+            module_list.sort(key=lambda named_module: named_module[1].weight.ds_tensor.numel()
+                             if is_zero3_enabled else named_module[1].weight.numel())
+        module_list = deque(module_list)
+
+        while len(module_list) > 0:
+            # Use popleft to timely release module's memory of replaced module after each loop iteration
+            module_name, module = module_list.popleft()
+
+            matched_key = None
+            matched_quantization_config = None
+
+            for key, config in quantization_config.items():
+                if key in module_name:
+                    assert matched_key is None, f'{module_name} matched multiple quantization key word {matched_key} and {key}'
+                    matched_key = key
+                    matched_quantization_config = config
+
+            if matched_key is None:
+                continue
+
+            if is_zero3_enabled:
+                module.weight.all_gather()
+
+            assert module.weight.dtype == torch.float16, 'Model weight is expected in half.'
+
+            new_module = QUANTIZATION_LAYER_MAPPINGS[type(module)](matched_quantization_config, module)
+
+            if is_zero3_enabled:
+                module.weight.partition()
+
+            recursive_setattr(model, module_name, new_module)
+
+            if matched_key not in matched_module_list_by_key:
+                matched_module_list_by_key[matched_key] = []
+            matched_module_list_by_key[matched_key].append(module_name)
+            matched_module_count += 1
+
+            # Timely recycle memory to prevent OOM on large models
+            gc.collect()
+
+    # Clear registry after model construction.
+    layers.quantized_weight_registry.clear()
+
+    logger.info(
+        f'Group-wise weight quantization summary: convert {matched_module_count} node(s) to quantized implementation')
+    summary_str = '\n'
+
+    for key, module_list in matched_module_list_by_key.items():
+        summary_str += f'Key: {key}, matched modules:\n'
+        for module_name in module_list:
+            summary_str += f'\t{module_name}\n'
+    logger.info(summary_str)
+
+    return model

deepspeed/inference/quantization/quantization_context.py

@@ -0,0 +1,13 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from deepspeed.runtime.zero import partition_parameters
+from deepspeed.runtime.swap_tensor.partitioned_param_swapper import AsyncPartitionedParameterSwapper
+
+
+class QuantizationContext(partition_parameters.Init):
+
+    def __init__(self, config_dict_or_path, param_swapper: AsyncPartitionedParameterSwapper = None) -> None:
+        super().__init__(config_dict_or_path=config_dict_or_path, param_swapper=param_swapper)

deepspeed/inference/quantization/utils.py

@@ -0,0 +1,288 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import deepspeed
+from torch import Tensor
+from typing import Tuple
+import torch.nn as nn
+from typing import Dict, Callable, Union
+from deepspeed.accelerator import get_accelerator
+import functools
+
+device = get_accelerator().device_name() if get_accelerator().is_available() else 'cpu'
+
+quantizer_module = None
+
+
+def get_quantizer_module():
+    global quantizer_module
+    if quantizer_module is None:
+        quantizer_module = deepspeed.ops.op_builder.QuantizerBuilder().load()
+    return quantizer_module
+
+
+def tensor_clamp(tensor: Tensor, min, max) -> Tensor:
+    if tensor.device.type == 'cpu' and tensor.dtype == torch.float16:
+        # CPU does not support FP16 clamp
+        return tensor.to(dtype=torch.float32).clamp_(min, max).to(dtype=torch.float16)
+    else:
+        return tensor.clamp_(min, max)
+
+
+def tensor_round(tensor: Tensor) -> Tensor:
+    if tensor.device.type == 'cpu' and tensor.dtype == torch.float16:
+        # CPU does not support FP16 round
+        return tensor.to(dtype=torch.float32).round_().to(dtype=torch.float16)
+    else:
+        return tensor.round_()
+
+
+class Quantizer:
+
+    def __init__(self, config: Dict) -> None:
+        self.config = config
+        assert self.config['num_bits'] == 4 or self.config[
+            'num_bits'] == 8, 'Only INT4 and INT8 quantization is supported.'
+        assert self.config['symmetric'] == False, 'Only asymmetric quantization is supported at this moment.'
+
+    def quantize(self, tensor: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+        assert tensor.shape[self.config['group_dim']] % self.config['group_size'] == 0 \
+            , f'Tensor shape: {tensor.shape} quantization config {self.config}'
+
+        tensor = torch.clone(tensor)
+
+        shape = tensor.shape
+        num_groups = shape[self.config['group_dim']] // self.config['group_size']
+        new_shape = (shape[:self.config['group_dim']] + (num_groups, self.config['group_size']) +
+                     shape[self.config['group_dim'] + 1:])
+        tensor = tensor.view(new_shape)
+
+        quantized_tensor, scale, min_value = self._quantize_int8(tensor)
+        quantized_tensor = quantized_tensor.view(shape)
+
+        if self.config['num_bits'] == 4:
+            return self._compress_uint8_to_uint4(quantized_tensor), scale, min_value
+        if self.config['num_bits'] == 8:
+            return quantized_tensor, scale, min_value
+
+        assert False, 'Unsupported quantization bits {}'.format(self.config['num_bits'])
+
+    def _quantize_int8(self, tensor: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+        q_range = 2**self.config['num_bits'] - 1
+        min_value = tensor.amin(dim=self.config['group_dim'] + 1, keepdim=True)
+        max_value = tensor.amax(dim=self.config['group_dim'] + 1, keepdim=True)
+
+        scale = q_range / (max_value - min_value)
+
+        tensor = tensor.sub_(min_value).mul_(scale)
+        tensor = tensor_round(tensor_clamp(tensor, 0, q_range)).to(torch.uint8)
+        return tensor, scale, min_value
+
+    def _compress_uint8_to_uint4(self, tensor: Tensor) -> Tensor:
+        assert tensor.shape[-1] % 2 == 0
+
+        new_data_shape = list(tensor.shape)
+        new_data_shape[-1] = new_data_shape[-1] // 2
+
+        data = torch.empty(new_data_shape, dtype=torch.uint8, device=tensor.device)
+        data = torch.bitwise_or(tensor[..., 0::2].bitwise_left_shift(4), tensor[..., 1::2])
+
+        return data
+
+
+class DeQuantizer:
+
+    def __init__(self, config: Dict, dtype: torch.dtype) -> None:
+        self.config = config
+        self.dtype = dtype
+        assert self.config['num_bits'] == 4 or self.config[
+            'num_bits'] == 8, 'Only INT4 and INT8 quantization is supported.'
+        assert self.config['symmetric'] == False, 'Only asymmetric quantization is supported at this moment.'
+
+    def dequantize(self, tensor: Tensor, quant_scale: Tensor, quant_min: Tensor) -> Tensor:
+        # Use customized CUDA quantization kernel if possible.
+        if self.config['group_size'] % 8 == 0 and \
+                (self.config['num_bits'] == 4 or self.config['num_bits'] == 8) and \
+                self.config['group_dim'] == len(tensor.shape) - 1 and \
+                    self.dtype == torch.float16 and device == get_accelerator().device_name():
+
+            last_dimension_size = self.config['group_size']
+            if self.config['num_bits'] == 4:
+                last_dimension_size = last_dimension_size // 2
+                quantized_tensor = get_quantizer_module().dequantize_int4_to_half_experimental(
+                    tensor.reshape(-1, last_dimension_size), quant_scale, quant_min,
+                    tensor.numel() // last_dimension_size, self.config['group_size'])
+                shape = list(tensor.shape)
+                shape[-1] = shape[-1] * 2
+            elif self.config['num_bits'] == 8:
+                # last_dimension_size = last_dimension_size // 2
+                quantized_tensor = get_quantizer_module().dequantize_int8_to_half_experimental(
+                    tensor.reshape(-1, last_dimension_size), quant_scale, quant_min,
+                    tensor.numel() // last_dimension_size, self.config['group_size'])
+                shape = list(tensor.shape)
+
+            return quantized_tensor.reshape(shape)
+
+        if self.config['num_bits'] == 4:
+            tensor = self._decompress_uint4_to_uint8(tensor)
+        elif self.config['num_bits'] != 8:
+            assert False, 'Unsupported quantization bits {}'.format(self.config['num_bits'])
+
+        shape = tensor.shape
+        num_groups = shape[self.config['group_dim']] // self.config['group_size']
+        new_shape = (shape[:self.config['group_dim']] + (num_groups, self.config['group_size']) +
+                     shape[self.config['group_dim'] + 1:])
+        tensor = tensor.view(new_shape)
+
+        dequantized_tensor = self._dequantize_int8(tensor, quant_scale, quant_min).view(shape)
+        return dequantized_tensor
+
+    def _dequantize_int8(self, tensor: Tensor, quant_scale: Tensor, quant_min: Tensor) -> Tensor:
+        assert tensor.dtype == torch.uint8
+        data = torch.zeros_like(tensor, dtype=self.dtype, device=tensor.device)
+        data = data.copy_(tensor)
+        data = data.div_(quant_scale).add_(quant_min)
+
+        return data
+
+    def _decompress_uint4_to_uint8(self, tensor: Tensor) -> Tensor:
+        new_data_shape = list(tensor.shape)
+        new_data_shape[-1] = new_data_shape[-1] * 2
+        data = torch.empty(new_data_shape, dtype=torch.uint8, device=tensor.device)
+        data[..., 0::2] = tensor.bitwise_right_shift(4)
+        data[..., 1::2] = tensor.bitwise_and(0xF)
+
+        return data
+
+
+def get_AsyncPartitionedParameterSwapper(model: nn.Module):
+    for param_name, param in model.named_parameters():
+        if hasattr(param, 'nvme_swapper') and param.nvme_swapper is not None:
+            return param.nvme_swapper
+    return None
+
+
+def recursive_setattr(model, module_name, module):
+    """
+    Recursively set the attribute of a module.
+    Args:
+        model (`torch.nn.Module`)
+            The model to set the attribute in.
+        module_name (`str`)
+            The name of the module to set the attribute in.
+        module (`torch.nn.Module`)
+            The module to set the attribute to.
+    """
+    split_list = module_name.split('.')
+    output = model
+    for name in split_list[:-1]:
+        output = getattr(output, name)
+    output.__setattr__(split_list[-1], module)
+
+
+def concat_to_compat_param(quantized_weight: Tensor,
+                           quant_scale: Tensor,
+                           quant_min: Tensor,
+                           return_param: bool = True) -> Union[nn.Parameter, Tensor]:
+    shape_wieght = quantized_weight.shape
+    shape_scale = quant_scale.shape
+    shape_min = quant_min.shape
+
+    quantized_weight = torch.flatten(quantized_weight)
+    quant_scale = torch.flatten(quant_scale)
+    quant_min = torch.flatten(quant_min)
+
+    def deconcat_individual_tensors(shape_wieght: torch.Size, shape_scale: torch.Size,
+                                    shape_min: torch.Size) -> Callable:
+
+        def fn(compat_tensor: nn.Parameter) -> Tuple[Tensor, Tensor, Tensor]:
+            weight = torch.narrow(compat_tensor, 0, 0, shape_wieght.numel()).view(shape_wieght)
+            scale = torch.narrow(compat_tensor, 0, shape_wieght.numel(), shape_scale.numel()).view(shape_scale)
+            min_val = torch.narrow(compat_tensor, 0,
+                                   shape_wieght.numel() + shape_scale.numel(), shape_min.numel()).view(shape_min)
+
+            return weight, scale, min_val
+
+        return fn
+
+    compat_tensor = torch.concat([quantized_weight, quant_scale, quant_min])
+    if return_param:
+        compat_tensor = nn.Parameter(compat_tensor, requires_grad=False)
+    compat_tensor.deconcat = deconcat_individual_tensors(shape_wieght, shape_scale, shape_min)
+
+    return compat_tensor
+
+
+def _quantize_param(param: nn.Parameter, quant_config: Dict):
+    assert not hasattr(param, 'weight_quantized'), 'Parameter has already been quantized.'
+    quantizer = Quantizer(quant_config)
+    dequantizer = DeQuantizer(quant_config, param.dtype)
+
+    quantized_weight, quant_scale, quant_min = quantizer.quantize(param.data)
+
+    quantized_weight = quantized_weight.view(param.dtype)
+    quant_scale = quant_scale.view(param.dtype)
+    quant_min = quant_min.view(param.dtype)
+
+    quantized_compat_tensor = concat_to_compat_param(quantized_weight, quant_scale, quant_min)
+    param.data = quantized_compat_tensor
+    param.deconcat = quantized_compat_tensor.deconcat
+
+    param.quantizer = quantizer
+    param.dequantizer = dequantizer
+    setattr(param, 'weight_quantized', True)
+
+
+def wrap_quantized_functional(f):
+
+    @functools.wraps(f)
+    def wrapper(input: Tensor, weight: nn.Parameter, *args, **kwargs) -> Tensor:
+        if hasattr(weight, 'weight_quantized') and getattr(weight, 'weight_quantized'):
+            quantized_weight, quant_scale, quant_min = weight.deconcat(weight)
+            temp_dequantized_weight = weight.dequantizer.dequantize(quantized_weight.view(torch.uint8), quant_scale,
+                                                                    quant_min)
+            return f(input, temp_dequantized_weight, *args, **kwargs)
+        else:
+            return f(input, weight, *args, **kwargs)
+
+    return wrapper
+
+
+def wrap_load_from_state_dict(f):
+
+    @functools.wraps(f)
+    def wrapper(model, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        replaced_old_value = None
+        key = None
+        # We may have nested wrappers if we launch multiple initialization context.
+        # Use state_dict_quantized flag to quantize state_dict only once
+        if hasattr(model.weight, 'weight_quantized') and getattr(
+                model.weight, 'weight_quantized') and not hasattr(model.weight, 'state_dict_quantized'):
+            setattr(model.weight, 'state_dict_quantized', True)
+            key = prefix + 'weight'
+            if key in state_dict:
+                quantized_weight, quant_scale, quant_min = model.weight.quantizer.quantize(state_dict[key])
+                quantized_weight = quantized_weight.view(model.weight.dtype)
+                quant_scale = quant_scale.view(model.weight.dtype)
+                quant_min = quant_min.view(model.weight.dtype)
+
+                replaced_old_value = state_dict[key]
+
+                state_dict[key] = concat_to_compat_param(quantized_weight, quant_scale, quant_min)
+
+        f(model, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+
+        if replaced_old_value is not None:
+            state_dict[key] = replaced_old_value
+            delattr(model.weight, 'state_dict_quantized')
+
+    return wrapper
+
+
+WEIGHT_QUANTIZATION_LAYERS = (
+    nn.Linear,
+    nn.Embedding,
+)

deepspeed/inference/v2/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+from .config_v2 import RaggedInferenceEngineConfig, DeepSpeedTPConfig
+from .engine_v2 import InferenceEngineV2
+from .engine_factory import build_hf_engine, build_engine_from_ds_checkpoint

deepspeed/inference/v2/allocator.py

@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from functools import reduce
+from typing import Iterable
+from collections import defaultdict
+import torch
+
+from deepspeed.accelerator import get_accelerator
+
+
+class Allocator:
+    cache = defaultdict(dict)
+
+    def empty_from(tensor: torch.Tensor, shape: Iterable[int]) -> torch.Tensor:
+        try:
+            return Allocator.cache[tensor][shape]
+        except KeyError:
+            shape_size = reduce(lambda x, y: x * y, shape)
+            if shape_size == 0:
+                raise ValueError("Cannot create empty tensor with size 0")
+            Allocator.cache[tensor][shape] = tensor.flatten()[:shape_size].view(shape)
+            return Allocator.cache[tensor][shape]
+
+
+empty_from = Allocator.empty_from
+
+
+def on_device(method) -> torch.Tensor:
+    """
+    Wraps a method to ensure the returned tensor is on the current device.
+    """
+
+    def wrapped(self, *args, **kwargs):
+        tensor = method(self, *args, **kwargs)
+        if isinstance(tensor, torch.Tensor):
+            return tensor.to(get_accelerator().current_device())
+        return tensor
+
+    return wrapped

deepspeed/inference/v2/checkpoint/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .base_engine import CheckpointEngineBase
+from .in_memory_engine import InMemoryModelEngine
+from .huggingface_engine import HuggingFaceCheckpointEngine

deepspeed/inference/v2/checkpoint/base_engine.py

@@ -0,0 +1,41 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import ABC, abstractmethod
+from typing import Iterable, Tuple
+
+import torch
+
+#from .huggingface_engine import HuggingFaceCheckpointEngine
+
+MEGATRON = 'megatron'
+HUGGINGFACE = 'huggingface'
+
+
+class CheckpointEngineBase(ABC):
+    """
+    Abstract interface for checkpoint engines to implement.
+
+    There is no ``__init__`` method here by design, since the creation of the checkpoint
+    engine will happen outside the policy/engine code. The tradeoff being made here is
+    that we will write different frontends for different checkpoint engines, but these
+    frontends can be tailored to the specific checkpoint engine/model source needs.
+    """
+
+    @abstractmethod
+    def parameters(self) -> Iterable[Tuple[str, torch.Tensor]]:
+        """
+        This method should create a generator of tuples of the form (name, parameter) for
+        all parameters in the model. The name should be the fully qualified name of the
+        parameter, and the parameter should be a torch.Tensor.
+
+        The expected use of a checkpoint engine is the following:
+        ```python
+        for name, parameter in checkpoint_engine.parameters():
+            container_map.map_param(name, parameter)
+        ```
+        For a concrete use example, see ``InferenceV2Policy``.
+        """
+        ...

deepspeed/inference/v2/checkpoint/huggingface_engine.py

@@ -0,0 +1,130 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import json
+import torch
+from .base_engine import CheckpointEngineBase
+from typing import Iterable, Tuple
+from functools import partial
+
+from ..logging import inference_logger
+
+
+class HuggingFaceCheckpointEngine(CheckpointEngineBase):
+
+    def __init__(self, model_name_or_path: str, auth_token: str = None, **hf_kwargs) -> None:
+        super().__init__()
+        from transformers import AutoConfig, GenerationConfig
+
+        self.model_name_or_path = model_name_or_path
+        self.auth_token = auth_token
+        self.model_config = AutoConfig.from_pretrained(self.model_name_or_path, **hf_kwargs)
+        # Define this property here so we can use it in the model implementation
+        if not hasattr(self.model_config, "max_seq_length"):
+            if hasattr(self.model_config, "max_position_embeddings"):
+                self.model_config.max_seq_length = self.model_config.max_position_embeddings
+            else:
+                generation_config = GenerationConfig.from_pretrained(self.model_name_or_path)
+                self.model_config.max_seq_length = generation_config.max_length
+        self._local_checkpoint_dir = None
+        self._all_ckpt_paths = self._fetch_checkpoint_files()
+
+    def _fetch_checkpoint_files(self):
+        """
+        Fetch the checkpoint files from the HuggingFace Hub.
+        """
+        # TODO(jeff): for models like llama-2 the user will have to provide an auth `token`,
+        # currently coming from the ckpt engine init but maybe a catch all kwargs for other
+        # snapshot download parameters would be more flexible.
+
+        from huggingface_hub import snapshot_download, list_repo_tree
+
+        def model_has_safetensors(model_name_or_path: str) -> bool:
+            if os.path.isdir(model_name_or_path):
+                file_list = os.listdir(model_name_or_path)
+            else:
+                file_list = [rf.path for rf in list_repo_tree(model_name_or_path)]
+            for f in file_list:
+                if f.endswith(".safetensors"):
+                    return True
+            return False
+
+        if os.path.isdir(self.model_name_or_path):
+            self._local_checkpoint_dir = self.model_name_or_path
+        else:
+            # We need to download the checkpoint files from HF
+            if model_has_safetensors(self.model_name_or_path):
+                # Prioritize downloading safetensors if they are available
+                allow_patterns = ["*.safetensors", "*.json"]
+            else:
+                # Fallback to bin files when safetensors are not present
+                allow_patterns = ["*.bin", "*.json", "*.pt"]
+            self._local_checkpoint_dir = snapshot_download(self.model_name_or_path,
+                                                           allow_patterns=allow_patterns,
+                                                           revision=None,
+                                                           token=self.auth_token)
+
+        assert os.path.isdir(
+            self._local_checkpoint_dir
+        ), f"Checkpoint dir {self._local_checkpoint_dir} is not a directory, cannot load checkpoint."
+
+        # Set the appropriate file names based on whether we have safetensors or not
+        if model_has_safetensors(self._local_checkpoint_dir):
+            from safetensors.torch import load_file
+            model_param_json_fname = "model.safetensors.index.json"
+            model_file_fname = "model.safetensors"
+            self._checkpoint_load_fn = load_file
+        else:
+            model_param_json_fname = "pytorch_model.bin.index.json"
+            model_file_fname = "pytorch_model.bin"
+            self._checkpoint_load_fn = partial(torch.load, map_location="cpu", weights_only=False)
+
+        model_param_json = os.path.join(self._local_checkpoint_dir, model_param_json_fname)
+
+        if not os.path.isfile(model_param_json):
+            # We don't need any json as all such HF models will have pytorch_model.bin
+            all_checkpoint_files = [os.path.join(self._local_checkpoint_dir, model_file_fname)]
+        else:
+            param_map = json.load(open(model_param_json, "r"))
+
+            # weight_map -> { "lm_head.weight": "pytorch_model-00002-of-00002.bin", ... }
+            weight_map = param_map["weight_map"]
+
+            # unique set of all checkpoint files
+            all_checkpoint_files = set(weight_map.values())
+
+            # get absolute path of all unique checkpoint files
+            all_checkpoint_files = [os.path.join(self._local_checkpoint_dir, f) for f in all_checkpoint_files]
+
+        return all_checkpoint_files
+
+    def parameters(self) -> Iterable[Tuple[str, torch.Tensor]]:
+        """
+        Generator of model parameters (satisfies the CheckpointEngineBase interface).
+        """
+        for checkpoint in self._all_ckpt_paths:
+            inference_logger().info(f"Loading checkpoint: {checkpoint}")
+            checkpoint_sd = self._checkpoint_load_fn(checkpoint)
+
+            # If the model has tied embeddings, we need to make sure the lm_head weights are tied to the embeddings weights
+            if hasattr(self.model_config, "tie_word_embeddings") and self.model_config.tie_word_embeddings:
+                if self.model_config.model_type == "qwen2":
+                    checkpoint_sd["lm_head.weight"] = checkpoint_sd["model.embed_tokens.weight"]
+
+            param_keys = list(checkpoint_sd.keys())
+            for param_name in param_keys:
+                param = checkpoint_sd[param_name]
+                yield param_name, param
+
+            del checkpoint_sd
+
+
+if __name__ == "__main__":
+    # To test, add your auth_token here and run `python huggingface_engine.py`
+    engine = HuggingFaceCheckpointEngine(model_name_or_path="meta-llama/Llama-2-7b-hf",
+                                         auth_token="hf_xxxxxxxxxxxxxxxxx")
+    for name, param in engine.parameters():
+        print(name, param.shape)

deepspeed/inference/v2/checkpoint/in_memory_engine.py

@@ -0,0 +1,40 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Iterable, Tuple
+import torch
+
+from .base_engine import CheckpointEngineBase
+
+
+class InMemoryModelEngine(CheckpointEngineBase):
+    """
+    This "checkpoint" engine uses the existing interface to enable loading parameters into an
+    inference model from a model already instantiated in memory. In general, this is not the
+    recommended way to use the inference engine, and should only be used when absolutely necessary.
+
+    The primary limitation of this approach is that the model must be fully instantiated in memory.
+    In a tensor parallel scenario, this means that the model is either replicated many times in host
+    memory. Currently, it is also recommended to only use this approach for models held in host memory.
+
+    In order to free the memory held by this copy of the model, we delete the model in the first call
+    to `parameters`, so it is not safe to make this call twice.
+    """
+
+    def __init__(self, model: torch.nn.Module) -> None:
+        """
+        Create virtual checkpoint engine for the provided module.
+
+        Args:
+            model (torch.nn.Module): Model to load parameters from.
+        """
+        super().__init__()
+        self.model = model
+
+    def parameters(self) -> Iterable[Tuple[str, torch.Tensor]]:
+        for name, parameter in self.model.named_parameters():
+            yield name, parameter
+
+        del self.model

deepspeed/inference/v2/config_v2.py

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from pydantic import Field
+from typing import Optional
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from .ragged import DSStateManagerConfig
+
+
+class DeepSpeedTPConfig(DeepSpeedConfigModel):
+    """ Configure tensor parallelism settings """
+
+    tp_size: int = 1
+    """ Number of devices to split the model across using tensor parallelism. """
+
+
+class QuantizationConfig(DeepSpeedConfigModel):
+    """ Configure tensor parallelism settings """
+
+    quantization_mode: Optional[str] = None
+    """ The quantization mode in string format. The supported modes are as follows:
+        - 'wf6af16', weight-only quantization with FP6 weight and FP16 activation.
+    """
+    # TODO: may reuse the constants in deepspeed/compression/constants.py
+
+
+class RaggedInferenceEngineConfig(DeepSpeedConfigModel):
+    """ Sets parameters for DeepSpeed Inference Engine. """
+
+    tensor_parallel: DeepSpeedTPConfig = Field({}, alias="tp")
+    """
+    Configuration for tensor parallelism used to split the model across several
+    GPUs. Expects a dictionary containing values for :any:`DeepSpeedTPConfig`.
+    """
+
+    state_manager: DSStateManagerConfig = Field({}, alias="manager")
+    """
+    Configuration for managing persistent state
+    """
+
+    quantization: QuantizationConfig = {}

deepspeed/inference/v2/engine_factory.py

@@ -0,0 +1,141 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import json
+import logging
+import os
+import pickle
+from packaging import version
+
+from .engine_v2 import InferenceEngineV2
+from .config_v2 import RaggedInferenceEngineConfig
+from .checkpoint import HuggingFaceCheckpointEngine
+from .logging import inference_logger
+from .model_implementations import (
+    Exaone4Policy,
+    OPTPolicy,
+    Llama2Policy,
+    MistralPolicy,
+    MixtralPolicy,
+    FalconPolicy,
+    PhiPolicy,
+    Phi3Policy,
+    QwenPolicy,
+    Qwen2Policy,
+    Qwen2MoePolicy,
+)
+from .model_implementations.inference_policy_base import POLICIES, InferenceV2Policy
+from .model_implementations.flat_model_helpers import make_metadata_filename, ModelMetadata
+
+
+def build_engine_from_ds_checkpoint(path: str,
+                                    engine_config: RaggedInferenceEngineConfig,
+                                    debug_level: int = logging.INFO) -> InferenceEngineV2:
+    """
+    Creates an engine from a checkpoint saved by ``InferenceEngineV2``.
+
+    Arguments:
+        path: Path to the checkpoint. This does not need to point to any files in particular,
+            just the directory containing the checkpoint.
+        engine_config: Engine configuration. See ``RaggedInferenceEngineConfig`` for details.
+        debug_level: Logging level to use. Unless you are actively seeing issues, the recommended
+            value is ``logging.INFO``.
+
+    Returns:
+        Fully initialized inference engine ready to serve queries.
+    """
+
+    inference_logger(level=debug_level)
+    # Load metadata, for grabbing the policy name we'll have all ranks just check for
+    # rank 0.
+    metadata_filename = make_metadata_filename(path, 0, engine_config.tensor_parallel.tp_size)
+    metadata = json.load(open(metadata_filename, "r"))
+    metadata = ModelMetadata.parse_raw(metadata)
+
+    # Get the policy
+    try:
+        policy_cls: InferenceV2Policy = POLICIES[metadata.policy]
+    except KeyError:
+        raise ValueError(f"Unknown policy {metadata.policy} for model {path}")
+
+    # Load the model config
+    model_config = pickle.load(open(os.path.join(path, "ds_model_config.pkl"), "rb"))
+    policy = policy_cls(model_config, inf_checkpoint_path=path)
+
+    return InferenceEngineV2(policy, engine_config)
+
+
+def build_hf_engine(path: str,
+                    engine_config: RaggedInferenceEngineConfig,
+                    debug_level: int = logging.INFO) -> InferenceEngineV2:
+    """
+    Build an InferenceV2 engine for HuggingFace models. This can accept both a HuggingFace
+    model name or a path to an Inference-V2 checkpoint.
+
+    Arguments:
+        path: Path to the checkpoint. This does not need to point to any files in particular,
+            just the directory containing the checkpoint.
+        engine_config: Engine configuration. See ``RaggedInferenceEngineConfig`` for details.
+        debug_level: Logging level to use. Unless you are actively seeing issues, the recommended
+            value is ``logging.INFO``.
+
+    Returns:
+        Fully initialized inference engine ready to serve queries.
+    """
+
+    if os.path.exists(os.path.join(path, "ds_model_config.pkl")):
+        return build_engine_from_ds_checkpoint(path, engine_config, debug_level=debug_level)
+    else:
+        # Set up logging
+        inference_logger(level=debug_level)
+        # get HF checkpoint engine
+        checkpoint_engine = HuggingFaceCheckpointEngine(path)
+
+        # get model config from HF AutoConfig
+        model_config = checkpoint_engine.model_config
+
+        # get the policy
+        # TODO: generalize this to other models
+        if model_config.model_type == "opt":
+            if not model_config.do_layer_norm_before:
+                raise ValueError(
+                    "Detected OPT-350m model. This model is not currently supported. If this is not the 350m model, please open an issue: https://github.com/deepspeedai/DeepSpeed-MII/issues"
+                )
+            policy = OPTPolicy(model_config, checkpoint_engine=checkpoint_engine)
+        elif model_config.model_type == "llama":
+            policy = Llama2Policy(model_config, checkpoint_engine=checkpoint_engine)
+        elif model_config.model_type == "mistral":
+            # Ensure we're using the correct version of transformers for mistral
+            import transformers
+            assert version.parse(transformers.__version__) >= version.parse("4.34.0"), \
+                f"Mistral requires transformers >= 4.34.0, you have version {transformers.__version__}"
+            policy = MistralPolicy(model_config, checkpoint_engine=checkpoint_engine)
+        elif model_config.model_type == "mixtral":
+            # Ensure we're using the correct version of transformers for mistral
+            import transformers
+            assert version.parse(transformers.__version__) >= version.parse("4.36.1"), \
+                f"Mistral requires transformers >= 4.36.1, you have version {transformers.__version__}"
+            policy = MixtralPolicy(model_config, checkpoint_engine=checkpoint_engine)
+        elif model_config.model_type == "falcon":
+            policy = FalconPolicy(model_config, checkpoint_engine=checkpoint_engine)
+        elif model_config.model_type == "phi":
+            policy = PhiPolicy(model_config, checkpoint_engine=checkpoint_engine)
+        elif model_config.model_type == "phi3":
+            policy = Phi3Policy(model_config, checkpoint_engine=checkpoint_engine)
+        elif model_config.model_type == "qwen":
+            policy = QwenPolicy(model_config, checkpoint_engine=checkpoint_engine)
+        elif model_config.model_type == "qwen2":
+            policy = Qwen2Policy(model_config, checkpoint_engine=checkpoint_engine)
+        elif model_config.model_type == "qwen2_moe":
+            policy = Qwen2MoePolicy(model_config, checkpoint_engine=checkpoint_engine)
+        elif model_config.model_type == "exaone4":
+            import transformers
+            assert version.parse(transformers.__version__) >= version.parse("4.54.0"), \
+                f"EXAONE 4.0 requires transformers >= 4.54.0, you have version {transformers.__version__}"
+            policy = Exaone4Policy(model_config, checkpoint_engine=checkpoint_engine)
+        else:
+            raise ValueError(f"Unsupported model type {model_config.model_type}")
+
+        return InferenceEngineV2(policy, engine_config)

deepspeed/inference/v2/engine_v2.py

@@ -0,0 +1,268 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import json
+import pickle
+from typing import Iterable, Tuple
+
+import torch
+
+import deepspeed.comm as dist
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.comm.comm import init_distributed
+
+from .model_implementations import InferenceV2Policy
+from .logging import inference_logger
+from .ragged import DSStateManager, RaggedBatchWrapper, PlaceholderSequenceDescriptor
+from .scheduling_utils import SchedulingError, SchedulingResult
+from .model_implementations.flat_model_helpers import make_param_filename, make_metadata_filename
+from .model_implementations.inference_model_base import DSInferenceModelBase
+
+from .config_v2 import RaggedInferenceEngineConfig
+
+INFERENCE_MODEL_TIMER = "model-forward-inference"
+
+
+class InferenceEngineV2:
+
+    _config: RaggedInferenceEngineConfig
+    """
+    Configuration of the inference engine.
+    """
+
+    _model: DSInferenceModelBase
+    """
+    Inference model supporting ragged inference.
+    """
+
+    _state_manager: DSStateManager
+    """
+    Persistent state manager for sequences and KV-cache.
+    """
+
+    @property
+    def free_blocks(self) -> torch.Tensor:
+        """
+        Number of free KV blocks. This is a tensor of shape [n_kv_cache_groups] where each
+        element is the number of free blocks in the corresponding KV cache group.
+        """
+        return self._state_manager.free_blocks
+
+    @property
+    def n_kv_cache_groups(self) -> int:
+        """
+        Number of KV cache groups.
+        """
+        return self._state_manager.n_kv_cache_groups
+
+    def model(self) -> DSInferenceModelBase:
+        """
+        The model implementation.
+        """
+        return self._model
+
+    def __init__(self, policy: InferenceV2Policy, engine_config: RaggedInferenceEngineConfig) -> None:
+        """
+        Create the Inference V2 engine.
+
+        Arguments:
+            policy (InferenceV2Policy): Policy for the model implementation. This policy object
+                will be used to build the model and load the checkpoint associated with it.
+            engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine.
+        """
+        self._config = engine_config
+        self._policy = policy
+        self._base_mp_group = self._initialize_tp_group()
+
+        # Build model from policy
+        inference_logger().info("Building model...")
+        self._model = self._policy.build_model(self._config, self._base_mp_group)
+        inference_logger().info("Model built.")
+
+        # Create state manager
+        self._batch = RaggedBatchWrapper(self._config.state_manager)
+        self._state_manager = DSStateManager(self._config.state_manager,
+                                             self._model.kv_cache_config(),
+                                             base_mp_group=self._base_mp_group)
+        self._model.set_state_manager(self._state_manager)
+
+    def _initialize_tp_group(self):
+        """
+        Implementation of our TP group initialization.
+        """
+        init_distributed()
+        local_rank = int(os.getenv("LOCAL_RANK", 0))
+        get_accelerator().set_device(local_rank)
+
+        if local_rank >= self._config.tensor_parallel.tp_size:
+            raise RuntimeError("Local rank is greater than TP size, ensure that the TP config is correct.")
+
+        ranks = list(range(self._config.tensor_parallel.tp_size))
+        return dist.new_group(ranks=ranks)
+
+    def put(self,
+            batch_uids: Iterable[int],
+            batch_tokens: Iterable[torch.Tensor],
+            do_checks: bool = True) -> torch.Tensor:
+        """
+        Put a ragged batch onto the inference engine. This will perform one forward and return
+        a Tensor of the shape [len(batch_uids), *output_shape]. Logits for the non-final tokens
+        are not calculated.
+
+        Arguments:
+            batch_uids: Iterable of uids for the batch on the host
+            batch_tokens: Iterable of token tensors for the batch on the host
+            do_checks: Check schedulability when it is set to True. You can skip this check for better performance when it has already been completed.
+        """
+
+        if do_checks:
+            token_lens = [len(tokens) for tokens in batch_tokens]
+            schedule_check = self.can_schedule(batch_uids, token_lens)
+            if schedule_check != SchedulingResult.Success:
+                raise SchedulingError(schedule_check)
+
+        self._batch.clear()
+        for uid, tokens in zip(batch_uids, batch_tokens):
+
+            host_seq_desc = self._state_manager.get_or_create_sequence(uid)
+            self._model.maybe_allocate_kv(host_seq_desc, tokens.numel())
+            host_seq_desc.pre_forward(tokens.numel())
+
+            # We can disable checks since we already validated schedulability.
+            self._batch.insert_sequence(host_seq_desc, tokens, do_checks=do_checks)
+
+        # Send all metadata to the device
+        self._batch.finalize()
+
+        # Prep all data structures for the actual forward (in anticipation of CG in the future)
+        # and also to amortize some of the costs in a more straightforward way.
+        self._model.prepare_batch(self._batch)
+
+        # Model implementation will pick up in the forward.
+        logits = self._model.forward(self._batch)
+
+        # We return one set of logits per sequence in the batch (saves cost on unembedding)
+        assert logits.shape[0] == self._batch.current_sequences
+
+        for uid in batch_uids:
+            host_seq_desc = self._state_manager.get_sequence(uid)
+            host_seq_desc.post_forward()  # Updates sequence metadata.
+            self._model.maybe_free_kv(host_seq_desc)
+
+        return logits
+
+    def query(self, uid: int, max_request_tokens: int, max_request_blocks) -> Tuple[int, torch.Tensor]:
+        """
+        Determine the number of tokens and KV blocks to reserve for a given request. Given a UID
+        (this UID may not be recognized by the model yet), this will return the number of tokens
+        and blocks to reserve for the request.
+
+        Arguments:
+            uid (int): The UID of the sequence (as tracked by the scheduling entity). If
+                this is a new sequence (with a UID unknown to the inference engine), then
+                an empty placeholder is created to pass to the occupancy logic.
+            n_tokens (int): The number of tokens to hypothetically send.
+
+        Returns:
+            Tuple[int, Optional[int]]: Tuple of free kv blocks and the number of blocks
+                required to schedule the sequence.
+        """
+        seq_desc = self._state_manager.get_sequence(uid)
+        if seq_desc is None:
+            if (self._state_manager.n_tracked_sequences == self._config.state_manager.max_tracked_sequences):
+                return (0, 0)
+            seq_desc = PlaceholderSequenceDescriptor()
+
+        req_tokens, req_blocks = self._model.get_kv_requirements(seq_desc, max_request_tokens, max_request_blocks)
+
+        return (req_tokens, req_blocks)
+
+    def can_schedule(self, uids: Iterable[int], lengths: Iterable[int]) -> SchedulingResult:
+        """
+        Dry run a batch to determine if it can be scheduled. Placeholder sequences will be
+        created for any UIDs that are unknown to the inference engine.
+
+        Arguments:
+            uids (Iterable[int]): Iterable of UIDs for the batch
+            lengths (Iterable[int]): Iterable of lengths for each sequence of the batch. This lengths
+                corresponds to the number of tokens to send in the hypothetical forward; history
+                tokens will be determined via UID lookup and future tokens are disregarded.
+
+        Returns:
+            bool: True if the batch can be scheduled, False otherwise.
+        """
+
+        cur_seqs = self._state_manager.n_tracked_sequences
+        free_blocks = self._state_manager.free_blocks
+        req_blocks = 0
+        batch_len = 0
+
+        if len(uids) > self._config.state_manager.max_ragged_sequence_count:
+            # Can only compose a batch from a limited number of sequences
+            return SchedulingResult.BatchSequenceLimitExceeded
+
+        for uid, length in zip(uids, lengths):
+            seq_desc = self._state_manager.get_sequence(uid)
+            if seq_desc is None:
+                cur_seqs += 1
+                seq_desc = PlaceholderSequenceDescriptor()
+
+            sched_len, sched_blocks = self._model.get_kv_requirements(seq_desc, length, free_blocks)
+
+            if sched_len != length:
+                # We ran out of KV cache
+                return SchedulingResult.KVCacheLimitExceeded
+
+            batch_len += length
+            free_blocks -= sched_blocks
+
+        if cur_seqs > self._config.state_manager.max_tracked_sequences:
+            # Would run out of tracking metadata
+            return SchedulingResult.EngineSequenceLimitExceeded
+
+        if batch_len > self._config.state_manager.max_ragged_batch_size:
+            # Would exceed the maximum batch size
+            return SchedulingResult.BatchTokenLimitExceeded
+
+        return SchedulingResult.Success
+
+    def get_remaining_block_capacity(self, uid: int) -> int:
+        """
+        Get the remaining capacity of the last block already allocated.
+        """
+        seq_desc = self._state_manager.get_sequence(uid)
+        if seq_desc is None:
+            return 0
+        return self._model.get_remaining_block_capacity(seq_desc)
+
+    def flush(self, uid: int) -> None:
+        """
+        Remove all state associated with a sequence from the inference engine.
+
+        Arguments:
+            uid (int): The UID of the sequence to flush.
+        """
+        self._state_manager.flush_sequence(uid)
+
+    def serialize(self, save_path: str) -> None:
+        """
+        Serialize the model to a file.
+
+        Arguments:
+            path (str): Path to the file to serialize to.
+        """
+        param_file_name = make_param_filename(save_path, self._model.tp_rank, self._model.tp_size)
+        metadata_file_name = make_metadata_filename(save_path, self._model.tp_rank, self._model.tp_size)
+
+        # Save the flattened parameters
+
+        torch.save(self._model.flattened_params, param_file_name)
+
+        json.dump(self._model.flattened_param_metadata.json(), open(metadata_file_name, "w"))
+
+        if self._model.tp_rank == 0:
+            pickle.dump(self._model._config, open(os.path.join(save_path, "ds_model_config.pkl"), "wb"))

deepspeed/inference/v2/inference_parameter.py

@@ -0,0 +1,89 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Dict
+
+import torch
+
+CORE_PARAM = "_ds_core_param_key"
+
+STR_TO_DTYPE = {
+    "torch.float32": torch.float32,
+    "torch.float64": torch.float64,
+    "torch.float16": torch.float16,
+    "torch.bfloat16": torch.bfloat16,
+    "torch.int64": torch.int64,
+    "torch.int32": torch.int32,
+    "torch.int16": torch.int16,
+    "torch.int8": torch.int8,
+    "torch.uint8": torch.uint8,
+    "torch.bool": torch.bool,
+}
+
+
+class InferenceParameter(torch.Tensor):
+    """
+    An extension of the torch.Tensor class to support our inference focused features. One important
+    thing to note here is that an InferenceParam can be used a torch.Tensor, but outputs of
+    torch.Tensor operations will not be InferenceParams.
+    """
+
+    @staticmethod
+    def __new__(cls, tensor, *args, **kwargs):
+        new_tensor = super().__new__(cls, tensor, *args, **kwargs)
+        if hasattr(tensor, "_aux_attrs"):
+            setattr(new_tensor, "_aux_attrs", tensor.aux_attrs)
+        return new_tensor
+
+    def to(self, *args, **kwargs):
+        new_tensor = super().to(*args, **kwargs)
+        if hasattr(self, "_aux_attrs"):
+            setattr(new_tensor, "_aux_attrs", self.aux_attrs)
+        try:
+            _ = torch.device(args[0])
+            for name, attr in new_tensor.aux_attrs.items():
+                new_attr = attr.to(*args, **kwargs)
+                setattr(new_tensor, name, new_attr)
+                new_tensor.aux_attrs[name] = new_attr
+        except Exception:
+            pass
+
+        return new_tensor
+
+    @classmethod
+    def initialize(cls, core_param: torch.Tensor, **kwargs) -> 'InferenceParameter':
+        """
+        Create the inference parameter.
+        """
+        param = InferenceParameter(core_param)
+        setattr(param, "_aux_attrs", kwargs)
+
+        for attr_name, attr in kwargs.items():
+            if hasattr(param, attr_name):
+                raise ValueError(f"Attribute {attr_name} already exists on param.")
+
+            if not isinstance(attr, torch.Tensor):
+                raise ValueError(f"Attribute {attr_name} must be a tensor.")
+
+            setattr(param, attr_name, attr)
+
+        return param
+
+    @classmethod
+    def initialize_raw(self, **kwargs) -> 'InferenceParameter':
+        """
+        All kwargs must be torch.Tensors and must include the core parameter.
+        """
+        if CORE_PARAM not in kwargs:
+            raise ValueError(f"Must provide core parameter, with key {CORE_PARAM}.")
+
+        return InferenceParameter.initialize(kwargs[CORE_PARAM], **kwargs)
+
+    @property
+    def aux_attrs(self) -> Dict[str, torch.Tensor]:
+        """
+        Dictionary of auxiliary attributes.
+        """
+        return self._aux_attrs

deepspeed/inference/v2/inference_utils.py

@@ -0,0 +1,105 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Dict
+
+import torch
+
+from enum import Enum, IntEnum
+
+
+class NormTypeEnum(Enum):
+    LayerNorm: str = "layer_norm"
+    RMSNorm: str = "rms_norm"
+
+
+class DtypeEnum(Enum):
+    # The torch dtype must always be the first value (so we return torch.dtype)
+    fp16 = torch.float16, "torch.float16", "fp16", "float16", "half"
+    fp32 = torch.float32, "torch.float32", "fp32", "float32", "float"
+    bf16 = torch.bfloat16, "torch.bfloat16", "bf16", "bfloat16", "bfloat"
+    int8 = torch.int8, "torch.int8", "int8"
+
+    # Copied from https://stackoverflow.com/a/43210118
+    # Allows us to use multiple values for each Enum index and returns first
+    # listed value when Enum is called
+    def __new__(cls, *values):
+        obj = object.__new__(cls)
+        # first value is canonical value
+        obj._value_ = values[0]
+        for other_value in values[1:]:
+            cls._value2member_map_[other_value] = obj
+        obj._all_values = values
+        return obj
+
+    def __repr__(self):
+        return "<%s.%s: %s>" % (
+            self.__class__.__name__,
+            self._name_,
+            ", ".join([repr(v) for v in self._all_values]),
+        )
+
+
+ELEM_SIZES: Dict[torch.dtype, int] = {
+    torch.float16: 2,
+    torch.bfloat16: 2,
+    torch.float32: 4,
+    torch.float64: 8,
+    torch.int8: 1,
+    torch.uint8: 1,
+    torch.int16: 2,
+    torch.int32: 4,
+    torch.int64: 8,
+    torch.bool: 1,
+}
+
+
+class ActivationType(IntEnum):
+    """
+    Types of activations supported by DS-Inference
+    """
+
+    GELU = 0
+
+    RELU = 1
+
+    SILU = 2
+
+    GEGLU = 3
+
+    ReGLU = 4
+
+    SiGLU = 5
+
+    IDENTITY = 6
+
+    InvalidType = -1
+
+
+def is_gated(act_fn: ActivationType) -> bool:
+    """
+    Return True if the given activation function is gated.
+    """
+    if not isinstance(act_fn, ActivationType):
+        act_fn = ActivationType(act_fn)
+
+    return act_fn in [ActivationType.GEGLU, ActivationType.ReGLU, ActivationType.SiGLU]
+
+
+def elem_size(dtype: torch.dtype) -> int:
+    """
+    Return size in bytes of the given dtype.
+    """
+    try:
+        return ELEM_SIZES[dtype]
+    except KeyError:
+        raise ValueError("Unknown dtype size for {}".format(dtype))
+
+
+def ceil_div(a: int, b: int) -> int:
+    """
+    Return ceil(a / b).
+    """
+    return -(-a // b)

deepspeed/inference/v2/kernels/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .ds_kernel import DSKernelBase

deepspeed/inference/v2/kernels/core_ops/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .bias_activations import *
+from .blas_kernels import *
+from .cuda_layer_norm import *
+from .cuda_rms_norm import *
+from .gated_activations import *
+from .cuda_linear import *

deepspeed/inference/v2/kernels/core_ops/bias_activations/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .bias_activation import *

deepspeed/inference/v2/kernels/core_ops/bias_activations/bias_activation.cpp

@@ -0,0 +1,68 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "bias_activation.h"
+#include <c10/cuda/CUDAStream.h>
+#include "ds_kernel_utils.h"
+
+#ifdef BF16_AVAILABLE
+#define DTYPE_SWITCH(DTYPE, ...)                                        \
+    [&] {                                                               \
+        if (DTYPE == torch::kFloat16) {                                 \
+            using scalar_t = __half;                                    \
+            return __VA_ARGS__();                                       \
+        } else if (DTYPE == torch::kBFloat16) {                         \
+            using scalar_t = __nv_bfloat16;                             \
+            return __VA_ARGS__();                                       \
+        } else {                                                        \
+            TORCH_CHECK(false, "Unsupported dtype for BiasActivation"); \
+        }                                                               \
+    }()
+#else
+#define DTYPE_SWITCH(DTYPE, ...)                                        \
+    [&] {                                                               \
+        if (DTYPE == torch::kFloat16) {                                 \
+            using scalar_t = __half;                                    \
+            return __VA_ARGS__();                                       \
+        } else {                                                        \
+            TORCH_CHECK(false, "Unsupported dtype for BiasActivation"); \
+        }                                                               \
+    }()
+#endif
+
+/*
+In-place bias and activation fusion kernel.
+*/
+void bias_activation(torch::Tensor& activation,
+                     c10::optional<torch::Tensor>& bias,
+                     const int32_t act_type)
+{
+    const ActivationType atype = static_cast<ActivationType>(act_type);
+    const int32_t rows = activation.size(0);
+    const int32_t cols = activation.size(1);
+
+    TORCH_CHECK(atype == ActivationType::GELU || atype == ActivationType::RELU ||
+                    atype == ActivationType::SILU || atype == ActivationType::IDENTITY,
+                "Unsupported activation type for BiasActivation");
+    TORCH_CHECK(activation.dim() == 2, "BiasActivation only supports 2D activation tensors");
+
+    DTYPE_SWITCH(activation.scalar_type(), [&] {
+        scalar_t* activation_ptr = reinterpret_cast<scalar_t*>(activation.data_ptr());
+
+        const scalar_t* bias_ptr;
+        if (bias.has_value()) {
+            TORCH_CHECK(activation.scalar_type() == bias.value().scalar_type(),
+                        "BiasActivation activation and bias must have same dtype");
+            bias_ptr = reinterpret_cast<const scalar_t*>(bias.value().data_ptr());
+        } else {
+            bias_ptr = nullptr;
+        }
+
+        if (atype == ActivationType::IDENTITY && bias_ptr == nullptr) { return; }
+
+        launch_bias_activation<scalar_t>(
+            activation_ptr, bias_ptr, rows, cols, atype, c10::cuda::getCurrentCUDAStream());
+    });
+}

deepspeed/inference/v2/kernels/core_ops/bias_activations/bias_activation.h

@@ -0,0 +1,22 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+#include "activation_type.h"
+
+template <typename T>
+void launch_bias_activation(T* activation,
+                            const T* bias,
+                            const int32_t n_rows,
+                            const int32_t n_cols,
+                            const ActivationType activation_type,
+                            cudaStream_t stream);
+
+void bias_activation(torch::Tensor& activation,
+                     c10::optional<torch::Tensor>& bias,
+                     const int32_t activation_type);

deepspeed/inference/v2/kernels/core_ops/bias_activations/bias_activation.py

@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional
+
+import torch
+
+from ....inference_utils import ActivationType, DtypeEnum
+from deepspeed.ops.op_builder import InferenceCoreBuilder
+from ... import DSKernelBase
+
+
+class CUDABiasActivation(DSKernelBase):
+    """
+    CUDA implementation of bias activation kernel. This kernel should be deprecated once
+    we are fusing the bias activation into the linear kernel in all scenarios.
+    """
+
+    supported_dtypes = [DtypeEnum.fp16, DtypeEnum.bf16]
+    supported_act_fns = [ActivationType.IDENTITY, ActivationType.GELU, ActivationType.RELU, ActivationType.SILU]
+
+    def __init__(self, channels: int, dtype: DtypeEnum, act_fn: ActivationType) -> None:
+        """
+        Compile and validate for the fused bias-activation kernel.
+
+        Parameters:
+            channels (int): Number of channels to expect in the activation.
+            dtype (torch.dtype): Data type for the input/output. Supported values
+                are DtypeEnum.fp16 and DtypeEnum.bf16.
+            act_fn (ActivationType): Activation function to use. Only IDENTITY, GELU, RELU, and SILU are supported.
+        """
+
+        if channels % 8 != 0:
+            raise ValueError("channels must be divisible by 8")
+
+        if DtypeEnum(dtype) not in CUDABiasActivation.supported_dtypes:
+            raise ValueError("Unsupported data type: {}, supported_dtypes are {}".format(
+                dtype, CUDABiasActivation.supported_dtypes))
+
+        act_fn = ActivationType(act_fn)
+        if act_fn not in CUDABiasActivation.supported_act_fns:
+            raise ValueError("Unsupported activation function: {}, supported_act_fns are {}".format(
+                act_fn, CUDABiasActivation.supported_act_fns))
+
+        inf_module = InferenceCoreBuilder().load()
+        self.kernel = inf_module.bias_activation
+        self.act_fn = act_fn
+
+    def __call__(self, activation: torch.Tensor, bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Add an optional bias and perform the non-linear activation function.
+
+        Parameters:
+            activation (torch.Tensor): Input tensor of shape [tokens, channels]
+            bias (torch.Tensor): Optional bias tensor of shape [channels]
+
+        Returns:
+            activation that has been updated in-place
+        """
+        self.kernel(activation, bias, self.act_fn.value)

deepspeed/inference/v2/kernels/core_ops/bias_activations/bias_activation_cuda.cu

@@ -0,0 +1,140 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include <cassert>
+#include "activation_type.h"
+#include "conversion_utils.h"
+#include "ds_kernel_utils.h"
+#include "memory_access_utils.h"
+
+// Default activation function will error out
+template <ActivationType ActType>
+DS_D_INLINE float act_fn(float val);
+
+template <>
+DS_D_INLINE float act_fn<ActivationType::IDENTITY>(float val)
+{
+    return val;
+}
+
+template <>
+DS_D_INLINE float act_fn<ActivationType::RELU>(float val)
+{
+    return val > 0.0f ? val : 0.0f;
+}
+
+template <>
+DS_D_INLINE float act_fn<ActivationType::GELU>(float val)
+{
+    constexpr float sqrt_param = 0.79788456080286535587989211986876f;
+    constexpr float mul_param = 0.044715f;
+    return val * 0.5f * (1.0f + tanhf(sqrt_param * (val + mul_param * val * val * val)));
+}
+
+template <>
+DS_D_INLINE float act_fn<ActivationType::SILU>(float val)
+{
+    return val / (1.0f + expf(-val));
+}
+
+namespace bias_act {
+
+constexpr int access_size = 16;
+constexpr int threads = 512;
+constexpr int unroll = 4;
+
+}  // namespace bias_act
+
+template <typename T, ActivationType ActType>
+__global__ void bias_activation_kernel(T* activation,
+                                       const T* bias,
+                                       const int32_t rows,
+                                       const int32_t cols)
+{
+    constexpr int vector_T = bias_act::access_size / sizeof(T);
+
+    const int32_t thread_offset = threadIdx.x * vector_T;
+    const int32_t block_offset = blockIdx.x * vector_T * bias_act::unroll * bias_act::threads;
+    const int32_t base_offset = block_offset + thread_offset;
+
+    const int32_t thread_stride = bias_act::threads * vector_T;
+
+#pragma unroll
+    for (int i = 0; i < bias_act::unroll; i++) {
+        const int32_t iter_offset = base_offset + i * thread_stride;
+
+        const int32_t row = iter_offset / cols;
+
+        T buffer[vector_T];
+        T bias_buffer[vector_T];
+
+        if (row < rows) {
+            const int32_t col = iter_offset % cols;
+
+            mem_access::load_global<bias_act::access_size>(buffer, activation + iter_offset);
+            mem_access::load_global<bias_act::access_size>(
+                bias_buffer, bias + col, bias != nullptr);
+
+#pragma unroll
+            for (int j = 0; j < vector_T; j++) {
+                float val =
+                    conversion::to<float>(buffer[j]) + conversion::to<float>(bias_buffer[j]);
+                buffer[j] = conversion::to<T>(act_fn<ActType>(val));
+            }
+
+            mem_access::store_global<bias_act::access_size>(activation + iter_offset, buffer);
+        }
+    }
+}
+
+#define ACT_TYPE_SWITCH(ACT_TYPE, ...)                                \
+    if (ACT_TYPE == ActivationType::IDENTITY) {                       \
+        constexpr ActivationType act_fn_t = ActivationType::IDENTITY; \
+        return __VA_ARGS__();                                         \
+    } else if (ACT_TYPE == ActivationType::RELU) {                    \
+        constexpr ActivationType act_fn_t = ActivationType::RELU;     \
+        return __VA_ARGS__();                                         \
+    } else if (ACT_TYPE == ActivationType::GELU) {                    \
+        constexpr ActivationType act_fn_t = ActivationType::GELU;     \
+        return __VA_ARGS__();                                         \
+    } else if (ACT_TYPE == ActivationType::SILU) {                    \
+        constexpr ActivationType act_fn_t = ActivationType::SILU;     \
+        return __VA_ARGS__();                                         \
+    } else {                                                          \
+        assert(false);                                                \
+    }
+
+template <typename T>
+void launch_bias_activation(T* activation,
+                            const T* bias,
+                            const int32_t n_rows,
+                            const int32_t n_cols,
+                            const ActivationType activation_type,
+                            cudaStream_t stream)
+{
+    constexpr int32_t elems_per_block =
+        bias_act::threads * bias_act::unroll * bias_act::access_size / sizeof(T);
+    const int32_t total_elems = n_rows * n_cols;
+
+    const int32_t blocks = (total_elems + elems_per_block - 1) / elems_per_block;
+
+    const dim3 grid(blocks);
+    const dim3 block(bias_act::threads);
+
+    ACT_TYPE_SWITCH(activation_type, [&] {
+        bias_activation_kernel<T, act_fn_t>
+            <<<grid, block, 0, stream>>>(activation, bias, n_rows, n_cols);
+    });
+}
+
+#define INSTANTIATE_FOR_T(T)                 \
+    template void launch_bias_activation<T>( \
+        T*, const T*, const int32_t, const int32_t, const ActivationType, cudaStream_t);
+
+INSTANTIATE_FOR_T(__half);
+
+#ifdef BF16_AVAILABLE
+INSTANTIATE_FOR_T(__nv_bfloat16);
+#endif

deepspeed/inference/v2/kernels/core_ops/blas_kernels/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .blas_linear import *

deepspeed/inference/v2/kernels/core_ops/blas_kernels/blas.h

@@ -0,0 +1,138 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+#include <cstdio>
+#include "blas_utils.h"
+
+#define DISPATCH_BLAS_MATMUL(T_TYPE, C_TYPE)                \
+    if (output.options().dtype() == torch::T_TYPE) {        \
+        blas_gemm_ex(output.data_ptr(),                     \
+                     (const void*)weights.data_ptr(),       \
+                     (const void*)hidden_states.data_ptr(), \
+                     m,                                     \
+                     n,                                     \
+                     k,                                     \
+                     lda,                                   \
+                     ldb,                                   \
+                     ldc,                                   \
+                     trans_a,                               \
+                     trans_b,                               \
+                     &alpha,                                \
+                     &beta,                                 \
+                     C_TYPE);                               \
+    }
+
+void blas_linear(at::Tensor& output, at::Tensor& hidden_states, at::Tensor& weights)
+{
+    /*
+    Expected shape: output([total_tokens_across_dims], out_neurons)
+                    hidden_states([total_tokens_across_dims], in_neurons)
+                    weights(out_neurons, in_neurons)
+
+    We are going to assume contiguous for the above shapes.
+
+    The shapes are going to get messed with a little internally to handle column-major
+    GEMMs.
+    */
+
+    // Number of tokens is N (since the GEMM output is column-major but our Tensor
+    // is row-major, we need to transpose the shapes)
+    const int n = output.numel() / output.size(-1);
+    const int k = weights.size(1);
+    const int m = weights.size(0);
+
+    // A strides
+    const bool trans_a = weights.stride(1) == 1;
+    const int lda = (trans_a) ? weights.stride(0) : weights.stride(1);
+
+    // B strides
+    const bool trans_b = hidden_states.stride(-1) != 1;
+    const int ldb = (trans_b) ? hidden_states.stride(-1) : hidden_states.stride(-2);
+
+    // C strides
+    const int ldc = output.stride(-2);
+
+    const float alpha = 1.0f;
+    const float beta = 0.0f;
+
+    TORCH_CHECK(output.scalar_type() == hidden_states.scalar_type(),
+                "Output and hidden states must have the same scalar type");
+    TORCH_CHECK(output.scalar_type() == weights.scalar_type(),
+                "Output and weights must have the same scalar type");
+
+    // Dispatch the datatypes
+    DISPATCH_BLAS_MATMUL(kFloat, BlasType::FP32);
+    DISPATCH_BLAS_MATMUL(kHalf, BlasType::FP16);
+#ifdef BF16_AVAILABLE
+    DISPATCH_BLAS_MATMUL(kBFloat16, BlasType::BF16);
+#endif
+}
+
+#define DISPATCH_4D_BLAS(T_TYPE, C_TYPE)                     \
+    if (C.options().dtype() == torch::T_TYPE) {              \
+        blas_strided_batched_gemm(C.data_ptr(),              \
+                                  (const void*)A.data_ptr(), \
+                                  (const void*)B.data_ptr(), \
+                                  m,                         \
+                                  n,                         \
+                                  k,                         \
+                                  lda,                       \
+                                  ldb,                       \
+                                  ldc,                       \
+                                  trans_a,                   \
+                                  trans_b,                   \
+                                  &alpha,                    \
+                                  &beta,                     \
+                                  stride_a,                  \
+                                  stride_b,                  \
+                                  stride_c,                  \
+                                  batch,                     \
+                                  C_TYPE);                   \
+    }
+
+void blas_4d_matmul(at::Tensor& C, at::Tensor& B, at::Tensor& A)
+{
+    /*
+    C shape: (batch_size, N, M)
+    A shape: (batch_size, N, K)
+    B shape: (batch_size, K, M)
+    */
+
+    const int n = C.size(-2);
+    const int k = C.size(-1);
+    const int m = B.size(-1);
+
+    // A strides
+    const bool trans_a = A.stride(-1) == 1;
+    const int lda = (trans_a) ? A.stride(-2) : A.stride(-1);
+    const int stride_a = A.stride(-3);
+
+    // B strides
+    const bool trans_b = B.stride(-1) != 1;
+    const int ldb = (trans_b) ? B.stride(-1) : B.stride(-2);
+    const int stride_b = B.stride(-3);
+
+    // C strides
+    const int ldc = C.stride(-2);
+    const int stride_c = C.stride(-3);
+
+    const float alpha = 1.0f;
+    const float beta = 0.0f;
+
+    const int batch = C.numel() / (n * m);
+
+    // Dispatch the datatypes
+    DISPATCH_4D_BLAS(kFloat, BlasType::FP32);
+    DISPATCH_4D_BLAS(kHalf, BlasType::FP16);
+#ifdef BF16_AVAILABLE
+    DISPATCH_4D_BLAS(kBFloat16, BlasType::BF16);
+#endif
+}
+
+void create_handle() { BlasContext::getInstance().get_handle(); }

deepspeed/inference/v2/kernels/core_ops/blas_kernels/blas_linear.py

@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ....inference_utils import DtypeEnum
+from deepspeed.ops.op_builder import InferenceCoreBuilder
+from ... import DSKernelBase
+
+
+class BlasLibLinear(DSKernelBase):
+    """
+    Wrapper around the BLAS matmul kernel for FP16/BF16/FP32 for CUDA/RoCM.
+
+    Performs z = x @ y
+    """
+
+    supported_dtypes = [DtypeEnum.fp16, DtypeEnum.bf16, DtypeEnum.fp32]
+
+    def __init__(self, fp_dtype: DtypeEnum):
+        """
+        Parameters:
+            fp_dtype (torch.dtype): Data type for the input/output. Supported values
+                are torch.float16, torch.bfloat16, and torch.float32.
+        """
+        fp_dtype = DtypeEnum(fp_dtype)
+        if fp_dtype not in BlasLibLinear.supported_dtypes:
+            raise ValueError("Unsupported data type: {}, supported_dtypes are {}".format(
+                fp_dtype, BlasLibLinear.supported_dtypes))
+
+        self.inf_module = InferenceCoreBuilder().load()
+        self.inf_module.create_handle()
+        self.kernel = self.inf_module.blas_linear
+
+    def __call__(self, output: torch.Tensor, hidden_states: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
+        """
+        Matmul kernel as implemented by platform BLAS library. The input must be 2D or larger. If
+        n-dimensional, the leading dimensions are folded into each other:
+            2D: m = x.size(0)
+            3D: m = x.size(0) * x.size(1)
+            4D: m = x.size(0) * x.size(1) * x.size(2) (etc...)
+        All inputs should be contiguous.
+
+        Parameters:
+            output (torch.Tensor): Output tensor. Shape is of [*, out_features]
+            hidden_states (torch.Tensor): Input tensor. Shape is of [*, in_features]
+            weights (torch.Tensor): Input tensor. Shape is of [out_features, in_features]
+
+        Returns:
+            z (torch.Tensor): Output tensor. Shape is of [m, n]
+        """
+        self.kernel(output, hidden_states, weights)
+        return output

deepspeed/inference/v2/kernels/core_ops/blas_kernels/blas_utils.h

@@ -0,0 +1,299 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <assert.h>
+#include <cublas_v2.h>
+#include <cuda.h>
+#ifdef BF16_AVAILABLE
+#include <cuda_bf16.h>
+#endif
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#ifndef __HIP_PLATFORM_AMD__
+#include <mma.h>
+#endif
+#include <stdio.h>
+#include <iostream>
+#include <stdexcept>
+
+class BlasContext {
+    /*
+    Slim wrapper for managing the lifetime of the platform's BLAS handle. This should
+    be hipified for ROCm.
+    */
+public:
+    BlasContext()
+    {
+        if (cublasCreate(&_handle) != CUBLAS_STATUS_SUCCESS) {
+            auto message = std::string("Fail to create cublas handle.");
+            std::cerr << message << std::endl;
+            throw std::runtime_error(message);
+        }
+#ifndef __HIP_PLATFORM_AMD__
+        cublasSetMathMode(_handle, CUBLAS_TENSOR_OP_MATH);
+#endif
+    }
+
+    virtual ~BlasContext() { cublasDestroy(_handle); }
+
+    static BlasContext& getInstance()
+    {
+        // Should always access the singleton through this function.
+        static BlasContext _instance;
+        return _instance;
+    }
+
+    cublasHandle_t get_handle() const { return _handle; }
+
+private:
+    cublasHandle_t _handle;
+};
+
+enum class BlasType { FP32, FP16, BF16 };
+
+// TODO HIP: Remove backward compatibility for torch<=2.0 in future
+#if defined(__HIP_PLATFORM_AMD__) && \
+    ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0))
+rocblas_operation get_trans_op(bool do_trans)
+{
+    return (do_trans) ? rocblas_operation_transpose : rocblas_operation_none;
+}
+
+rocblas_datatype get_datatype(BlasType type)
+{
+    switch (type) {
+        case BlasType::FP32: return rocblas_datatype_f32_r;
+        case BlasType::FP16: return rocblas_datatype_f16_r;
+        case BlasType::BF16: return rocblas_datatype_bf16_r;
+        default: throw std::runtime_error("Unsupported BlasType");
+    }
+}
+#else
+cublasOperation_t get_trans_op(bool do_trans) { return (do_trans) ? CUBLAS_OP_T : CUBLAS_OP_N; }
+
+cublasDataType_t get_datatype(BlasType type)
+{
+    switch (type) {
+#ifdef __HIP_PLATFORM_AMD__
+        case BlasType::FP32: return HIPBLAS_R_32F;
+        case BlasType::FP16: return HIPBLAS_R_16F;
+        case BlasType::BF16: return HIPBLAS_R_16B;
+#else
+        case BlasType::FP32: return CUDA_R_32F;
+        case BlasType::FP16: return CUDA_R_16F;
+        case BlasType::BF16: return CUDA_R_16BF;
+#endif
+        default: throw std::runtime_error("Unsupported BlasType");
+    }
+}
+#endif
+
+int blas_gemm_ex(void* C,
+                 const void* A,
+                 const void* B,
+                 int m,
+                 int n,
+                 int k,
+                 int lda,
+                 int ldb,
+                 int ldc,
+                 bool transa,
+                 bool transb,
+                 const float* alpha,
+                 const float* beta,
+                 BlasType type)
+{
+#if defined(__HIP_PLATFORM_AMD__) && \
+    ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0))
+    rocblas_operation_t transa_op = get_trans_op(transa);
+    rocblas_operation_t transb_op = get_trans_op(transb);
+
+    rocblas_datatype_t abc_type = get_datatype(type);
+
+    rocblas_status status = rocblas_gemm_ex(BlasContext::getInstance().get_handle(),
+                                            transa_op,
+                                            transb_op,
+                                            m,
+                                            n,
+                                            k,
+                                            (const void*)alpha,
+                                            A,
+                                            abc_type,
+                                            lda,
+                                            B,
+                                            abc_type,
+                                            ldb,
+                                            (const void*)beta,
+                                            C,
+                                            abc_type,
+                                            ldc,
+                                            C,
+                                            abc_type,
+                                            ldc,
+                                            rocblas_datatype_f32_r,
+                                            rocblas_gemm_algo_standard,
+                                            0,
+                                            0);
+#else
+    cublasOperation_t transa_op = get_trans_op(transa);
+    cublasOperation_t transb_op = get_trans_op(transb);
+
+    cublasDataType_t abc_type = get_datatype(type);
+    cublasStatus_t status = cublasGemmEx(BlasContext::getInstance().get_handle(),
+                                         transa_op,
+                                         transb_op,
+                                         m,
+                                         n,
+                                         k,
+                                         (const void*)alpha,
+                                         A,
+                                         abc_type,
+                                         lda,
+                                         B,
+                                         abc_type,
+                                         ldb,
+                                         (const void*)beta,
+                                         C,
+                                         abc_type,
+                                         ldc,
+#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2)
+                                         HIPBLAS_COMPUTE_32F,
+#elif defined(__HIP_PLATFORM_AMD__)
+                                         HIPBLAS_R_32F,
+#else
+                                         CUDA_R_32F,
+#endif
+                                         CUBLAS_GEMM_DEFAULT_TENSOR_OP);
+#endif
+
+#if defined(__HIP_PLATFORM_AMD__) && \
+    ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0))
+    if (status != rocblas_status_success) {
+#else
+    if (status != CUBLAS_STATUS_SUCCESS) {
+#endif
+        fprintf(stderr,
+                "!!!! kernel execution error. (m: %d, n: %d, k: %d, error: %d) \n",
+                m,
+                n,
+                k,
+                (int)status);
+        return EXIT_FAILURE;
+    }
+    return 0;
+}
+
+int blas_strided_batched_gemm(void* C,
+                              const void* A,
+                              const void* B,
+                              int m,
+                              int n,
+                              int k,
+                              int lda,
+                              int ldb,
+                              int ldc,
+                              bool transa,
+                              bool transb,
+                              const float* alpha,
+                              const float* beta,
+                              int stride_A,
+                              int stride_B,
+                              int stride_C,
+                              int batch,
+                              BlasType type)
+{
+#if defined(__HIP_PLATFORM_AMD__) && \
+    ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0))
+    rocblas_operation_t transa_op = get_trans_op(transa);
+    rocblas_operation_t transb_op = get_trans_op(transb);
+
+    rocblas_datatype_t abc_type = get_datatype(type);
+
+    rocblas_status status =
+        rocblas_gemm_strided_batched_ex(BlasContext::getInstance()::get_handle(),
+                                        transa_op,
+                                        transb_op,
+                                        m,
+                                        n,
+                                        k,
+                                        (const void*)alpha,
+                                        A,
+                                        abc_type,
+                                        lda,
+                                        stride_A,
+                                        B,
+                                        abc_type,
+                                        ldb,
+                                        stride_B,
+                                        (const void*)beta,
+                                        C,
+                                        abc_type,
+                                        ldc,
+                                        stride_C,
+                                        C,
+                                        abc_type,
+                                        ldc,
+                                        stride_C,
+                                        batch,
+                                        rocblas_datatype_f32_r,
+                                        rocblas_gemm_algo_standard,
+                                        0,
+                                        0);
+#else
+    cublasOperation_t transa_op = get_trans_op(transa);
+    cublasOperation_t transb_op = get_trans_op(transb);
+
+    cublasDataType_t abc_type = get_datatype(type);
+
+    cublasStatus_t status = cublasGemmStridedBatchedEx(BlasContext::getInstance().get_handle(),
+                                                       transa_op,
+                                                       transb_op,
+                                                       m,
+                                                       n,
+                                                       k,
+                                                       (const void*)alpha,
+                                                       A,
+                                                       abc_type,
+                                                       lda,
+                                                       stride_A,
+                                                       B,
+                                                       abc_type,
+                                                       ldb,
+                                                       stride_B,
+                                                       (const void*)beta,
+                                                       C,
+                                                       abc_type,
+                                                       ldc,
+                                                       stride_C,
+                                                       batch,
+#if defined(__HIP_PLATFORM_AMD__) && defined(HIPBLAS_V2)
+                                                       HIPBLAS_COMPUTE_32F,
+#elif defined(__HIP_PLATFORM_AMD__)
+                                                       HIPBLAS_R_32F,
+#else
+                                                       CUDA_R_32F,
+#endif
+                                                       CUBLAS_GEMM_DEFAULT_TENSOR_OP);
+#endif
+
+#if defined(__HIP_PLATFORM_AMD__) && \
+    ((TORCH_VERSION_MAJOR < 2) || (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR == 0))
+    if (status != rocblas_status_success) {
+#else
+    if (status != CUBLAS_STATUS_SUCCESS) {
+#endif
+        fprintf(stderr,
+                "!!!! kernel execution error. (batch: %d, m: %d, n: %d, k: %d, error: %d) \n",
+                batch,
+                m,
+                n,
+                k,
+                (int)status);
+        return EXIT_FAILURE;
+    }
+    return 0;
+}

deepspeed/inference/v2/kernels/core_ops/core_ops.cpp

@@ -0,0 +1,42 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+
+#include "bias_activation.h"
+#include "blas.h"
+#include "gated_activation_kernels.h"
+#include "layer_norm.h"
+#include "linear_kernels.h"
+#include "rms_norm.h"
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    // bias_activation.h
+    m.def("bias_activation", &bias_activation, "DeepSpeed bias activation in CUDA");
+
+    // layer_norm.h
+    m.def("layer_norm", &ds_layer_norm, "DeepSpeed layer norm in CUDA");
+    m.def("pre_layer_norm", &ds_pre_layer_norm, "DeepSpeed pre layer norm in CUDA");
+    m.def("post_layer_norm", &ds_post_layer_norm, "DeepSpeed pre layer norm in CUDA");
+
+    // blas.h
+    m.def("blas_linear", &blas_linear, "Linear implemented by vendor BLAS");
+    m.def("blas_4d_matmul", &blas_4d_matmul, "4D matmul implemented by vendor BLAS");
+    m.def("create_handle", &create_handle, "Create a handle for vendor BLAS");
+
+    // gated_activation_kernels.h
+    m.def("gated_activation", &ds_gated_activation, "DeepSpeed gated activation in CUDA");
+
+    // rms_norm.h
+    m.def("rms_norm", &rms_norm, "DeepSpeed rms norm in CUDA");
+    m.def("rms_pre_norm", &rms_pre_norm, "DeepSpeed rms pre norm in CUDA");
+
+    // linear_kernels.h
+    m.def("cuda_wf6af16_linear", &cuda_wf6af16_linear, "DeepSpeed Wf6Af16 linear in CUDA");
+    m.def(
+        "preprocess_weight", &preprocess_weight, "preprocess the FP16 weight to be 2bit and 4 bit");
+}

deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .cuda_ln import *
+from .cuda_post_ln import *
+from .cuda_pre_ln import *

deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/cuda_fp_ln_base.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ... import DSKernelBase
+from ....inference_utils import elem_size
+from deepspeed.ops.op_builder import InferenceCoreBuilder
+
+
+class CUDAFPLNBase(DSKernelBase):
+    """
+    Base class for CUDA LN kernels. They all same the same validation logic,
+    so we can share it here.
+    """
+
+    supported_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+
+    def __init__(self, channels: int, fp_dtype: torch.dtype, epsilon: float = 1e-5):
+        """
+        Parameters:
+            channels (int): Number of channels in the input tensor. Must be divisible to align
+                to 16 bytes.
+            fp_dtype (torch.dtype): Data type for the input/output/gamma. Supported values
+                are torch.float16, torch.bfloat16, and torch.float32.
+        """
+        if fp_dtype not in CUDAFPLNBase.supported_dtypes:
+            raise ValueError("Unsupported data type: {}, supported_dtypes are {}".format(
+                fp_dtype, CUDAFPLNBase.supported_dtypes))
+
+        if elem_size(fp_dtype) * channels % 16 != 0:
+            raise ValueError("channels must be divisible by 16 bytes")
+
+        self.inf_module = InferenceCoreBuilder().load()
+        self.epsilon = epsilon

deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/cuda_ln.py

@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from .cuda_fp_ln_base import CUDAFPLNBase
+
+
+class CUDAFPLN(CUDAFPLNBase):
+    """
+    Floating point layer norm kernel for CUDA/RoCM.
+
+    Performs: z = ln(x)
+    """
+
+    def __call__(self, output_z: torch.Tensor, input_x: torch.Tensor, gamma: torch.Tensor,
+                 beta: torch.Tensor) -> torch.Tensor:
+        """
+        output_z may alias input_x directly. All Tensors should have the same shape.
+
+        Parameters:
+            output_z (torch.Tensor): Output tensor.
+            input_x (torch.Tensor): Input tensor.
+            gamma (torch.Tensor): Gamma tensor.
+            beta (torch.Tensor): Beta tensor.
+        """
+        self.inf_module.layer_norm(output_z, input_x, gamma, beta, self.epsilon)
+        return output_z

deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/cuda_post_ln.py

@@ -0,0 +1,34 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from .cuda_fp_ln_base import CUDAFPLNBase
+
+
+class CUDAFPPostLN(CUDAFPLNBase):
+    """
+    Floating point post-LayerNorm kernel for CUDA/RoCM.
+
+    Performs: z = ln(x + y)
+    """
+
+    def __call__(self, output_z: torch.Tensor, input_x: torch.Tensor, input_y: torch.Tensor, gamma: torch.Tensor,
+                 beta: torch.Tensor) -> torch.Tensor:
+        """
+        Either input_x or input_y can alias output_z.
+
+        Parameters:
+            output_z (torch.Tensor): Output tensor.
+            input_x (torch.Tensor): Input tensor.
+            input_y (torch.Tensor): Input tensor.
+            gamma (torch.Tensor): Gamma tensor.
+            beta (torch.Tensor): Beta tensor.
+
+        Returns:
+            output (torch.Tensor): Output tensor.
+        """
+        self.inf_module.post_layer_norm(output_z, input_x, input_y, gamma, beta, self.epsilon)
+        return output_z

deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/cuda_pre_ln.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Tuple
+
+import torch
+
+from .cuda_fp_ln_base import CUDAFPLNBase
+
+
+class CUDAFPPreLN(CUDAFPLNBase):
+    """
+    Floating point pre-LayerNorm kernel for CUDA/RoCM.
+
+    Performs: z_res = x_res + y_hid
+              z_hid = ln(z_hid)
+    """
+
+    def __call__(self, z_res: torch.Tensor, z_hid: torch.Tensor, x_res: torch.Tensor, y_hid: torch.Tensor,
+                 gamma: torch.Tensor, beta: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        z_res can alias x_res. All non-parameter input/output tensors
+        must have the same shape. z_hid can alias y_hid.
+
+        Parameters:
+            z_res (torch.Tensor): Output residual.
+            z_hid (torch.Tensor): Output hidden states.
+            x_res (torch.Tensor): Input residual.
+            y_hid (torch.Tensor): Input hidden states.
+            gamma (torch.Tensor): Gamma tensor.
+            beta (torch.Tensor): Beta tensor.
+
+        Returns:
+            output (torch.Tensor): Output tensor.
+        """
+        self.inf_module.pre_layer_norm(z_res, z_hid, x_res, y_hid, gamma, beta, self.epsilon)
+        return z_res, z_hid

deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm.cpp

@@ -0,0 +1,102 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "layer_norm.h"
+
+#define DISPATCH_LAYER_NORM(T_TYPE, C_TYPE)                \
+    if (input.options().dtype() == torch::T_TYPE) {        \
+        launch_fused_ln((C_TYPE*)output.data_ptr(),        \
+                        (const C_TYPE*)input.data_ptr(),   \
+                        (const C_TYPE*)gamma.data_ptr(),   \
+                        (const C_TYPE*)beta.data_ptr(),    \
+                        epsilon,                           \
+                        rows,                              \
+                        elems_per_row,                     \
+                        at::cuda::getCurrentCUDAStream()); \
+    }
+
+void ds_layer_norm(at::Tensor& output,
+                   at::Tensor& input,
+                   at::Tensor& gamma,
+                   at::Tensor& beta,
+                   float epsilon)
+{
+    bool ragged_input = input.dim() == 2;
+
+    const int rows = ragged_input ? input.size(0) : input.size(0) * input.size(1);
+    const int elems_per_row = ragged_input ? input.size(1) : input.size(2);
+
+    DISPATCH_LAYER_NORM(kFloat, float);
+    DISPATCH_LAYER_NORM(kHalf, __half);
+#ifdef BF16_AVAILABLE
+    DISPATCH_LAYER_NORM(kBFloat16, __nv_bfloat16);
+#endif
+}
+
+#define DISPATCH_LAYER_NORM_RESIDUAL(T_TYPE, C_TYPE)             \
+    if (input.options().dtype() == torch::T_TYPE) {              \
+        launch_fused_post_ln((C_TYPE*)output.data_ptr(),         \
+                             (const C_TYPE*)input.data_ptr(),    \
+                             (const C_TYPE*)residual.data_ptr(), \
+                             (const C_TYPE*)gamma.data_ptr(),    \
+                             (const C_TYPE*)beta.data_ptr(),     \
+                             epsilon,                            \
+                             rows,                               \
+                             elems_per_row,                      \
+                             at::cuda::getCurrentCUDAStream());  \
+    }
+
+void ds_post_layer_norm(at::Tensor& output,
+                        at::Tensor& input,
+                        at::Tensor& residual,
+                        at::Tensor& gamma,
+                        at::Tensor& beta,
+                        float epsilon)
+{
+    bool ragged_input = input.dim() == 2;
+
+    const int rows = ragged_input ? input.size(0) : input.size(0) * input.size(1);
+    const int elems_per_row = ragged_input ? input.size(1) : input.size(2);
+
+    DISPATCH_LAYER_NORM_RESIDUAL(kFloat, float);
+    DISPATCH_LAYER_NORM_RESIDUAL(kHalf, __half);
+#ifdef BF16_AVAILABLE
+    DISPATCH_LAYER_NORM_RESIDUAL(kBFloat16, __nv_bfloat16);
+#endif
+}
+
+#define DISPATCH_PRE_LAYER_NORM_RESIDUAL(T_TYPE, C_TYPE)        \
+    if (input.options().dtype() == torch::T_TYPE) {             \
+        launch_fused_pre_ln((C_TYPE*)norm_output.data_ptr(),    \
+                            (C_TYPE*)res_output.data_ptr(),     \
+                            (const C_TYPE*)input.data_ptr(),    \
+                            (const C_TYPE*)residual.data_ptr(), \
+                            (const C_TYPE*)gamma.data_ptr(),    \
+                            (const C_TYPE*)beta.data_ptr(),     \
+                            epsilon,                            \
+                            rows,                               \
+                            elems_per_row,                      \
+                            at::cuda::getCurrentCUDAStream());  \
+    }
+
+void ds_pre_layer_norm(at::Tensor& res_output,
+                       at::Tensor& norm_output,
+                       at::Tensor& input,
+                       at::Tensor& residual,
+                       at::Tensor& gamma,
+                       at::Tensor& beta,
+                       float epsilon)
+{
+    bool ragged_input = input.dim() == 2;
+
+    const int rows = ragged_input ? input.size(0) : input.size(0) * input.size(1);
+    const int elems_per_row = ragged_input ? input.size(1) : input.size(2);
+
+    DISPATCH_PRE_LAYER_NORM_RESIDUAL(kFloat, float);
+    DISPATCH_PRE_LAYER_NORM_RESIDUAL(kHalf, __half);
+#ifdef BF16_AVAILABLE
+    DISPATCH_PRE_LAYER_NORM_RESIDUAL(kBFloat16, __nv_bfloat16);
+#endif
+}

deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm.h

@@ -0,0 +1,67 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+#include "ds_kernel_utils.h"
+
+/*
+Kernel launch methods for layer norm variants.
+*/
+
+template <typename T>
+void launch_fused_ln(T* output,
+                     const T* vals,
+                     const T* gamma,
+                     const T* beta,
+                     float epsilon,
+                     int rows,
+                     int elems_per_row,
+                     cudaStream_t stream);
+
+template <typename T>
+void launch_fused_post_ln(T* output,
+                          const T* vals,
+                          const T* residual,
+                          const T* gamma,
+                          const T* beta,
+                          float epsilon,
+                          int rows,
+                          int elems_per_row,
+                          cudaStream_t stream);
+template <typename T>
+void launch_fused_pre_ln(T* norm_output,
+                         T* res_output,
+                         const T* vals,
+                         const T* residual,
+                         const T* gamma,
+                         const T* beta,
+                         float epsilon,
+                         int rows,
+                         int elems_per_row,
+                         cudaStream_t stream);
+
+void ds_layer_norm(at::Tensor& output,
+                   at::Tensor& input,
+                   at::Tensor& gamma,
+                   at::Tensor& beta,
+                   float epsilon);
+
+void ds_post_layer_norm(at::Tensor& output,
+                        at::Tensor& input,
+                        at::Tensor& residual,
+                        at::Tensor& gamma,
+                        at::Tensor& beta,
+                        float epsilon);
+
+void ds_pre_layer_norm(at::Tensor& res_output,
+                       at::Tensor& norm_output,
+                       at::Tensor& input,
+                       at::Tensor& residual,
+                       at::Tensor& gamma,
+                       at::Tensor& beta,
+                       float epsilon);

deepspeed/inference/v2/kernels/core_ops/cuda_layer_norm/layer_norm_cuda.cu

@@ -0,0 +1,489 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "conversion_utils.h"
+#include "ds_kernel_utils.h"
+#include "memory_access_utils.h"
+#include "reduction_utils.h"
+
+namespace cg = cooperative_groups;
+using rop = reduce::ROpType;
+
+namespace ln {
+constexpr int granularity = 16;
+}  // namespace ln
+
+/*
+Regular layer norm implementation. Assumes elems_per_row % 8
+is equal to 0.
+
+Args:
+    output: buffer for output data
+    vals: buffer for input data
+    gamma: gain for normalization
+    beta: bias for normalization
+    epsilon: numeric stability
+    elems_per_row: number of elements each block will normalize
+*/
+template <typename T, int unRoll, int threadsPerGroup, int maxThreads>
+__global__ void fused_ln(T* output,
+                         const T* vals,
+                         const T* gamma,
+                         const T* beta,
+                         float epsilon,
+                         int elems_per_row)
+{
+    constexpr int T_per_load = ln::granularity / sizeof(T);
+
+    cg::thread_block tb = cg::this_thread_block();
+    cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);
+
+    // X-dimension of the block
+    const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) +
+                             (tb.thread_index().y * elems_per_row);
+    const int thread_offset = tb.thread_index().x * T_per_load;
+    const int base_offset = block_offset + thread_offset;
+    const int stride = blockDim.x * T_per_load;
+
+    float sum = reduce::init<rop::Add, float>();
+
+    const T* input_base = vals + base_offset;
+
+    T local_buffer[unRoll * T_per_load];
+
+#pragma unRoll
+    for (int i = 0; i < unRoll; i++) {
+        T* iteration_buffer = local_buffer + i * T_per_load;
+
+        mem_access::load_global<ln::granularity>(
+            iteration_buffer, input_base + i * stride, thread_offset + i * stride < elems_per_row);
+
+#pragma unRoll
+        for (int j = 0; j < T_per_load; j++) {
+            float vals_up_cast = conversion::to<float>(iteration_buffer[j]);
+            sum = reduce::element<rop::Add>(sum, vals_up_cast);
+        }
+    }
+
+    reduce::partitioned_block<rop::Add, threadsPerGroup>(tb, warp, sum);
+    const float mean = sum / elems_per_row;
+
+    float mean_diff = reduce::init<rop::Add, float>();
+
+#pragma unRoll
+    for (int i = 0; i < unRoll; i++) {
+#pragma unRoll
+        for (int j = 0; j < T_per_load; j++) {
+            // Using a 0 value here skews the variance, have to if-guard
+            if (thread_offset + i * stride < elems_per_row) {
+                float diff = (conversion::to<float>(local_buffer[i * T_per_load + j]) - mean);
+                mean_diff = reduce::element<rop::Add>(mean_diff, diff * diff);
+            }
+        }
+    }
+
+    reduce::partitioned_block<rop::Add, threadsPerGroup>(tb, warp, mean_diff);
+    const float variance = mean_diff / elems_per_row;
+    const float denom = __frsqrt_rn(variance + epsilon);
+
+    T* block_output = output + block_offset;
+
+#pragma unRoll
+    for (int i = 0; i < unRoll; i++) {
+        T* iteration_buffer = local_buffer + i * T_per_load;
+        const int iter_idx = i * stride + thread_offset;
+        const bool do_loads = iter_idx < elems_per_row;
+
+        T gamma_local[T_per_load], beta_local[T_per_load];
+
+        mem_access::load_global<ln::granularity>(gamma_local, gamma + iter_idx, do_loads);
+        mem_access::load_global<ln::granularity>(beta_local, beta + iter_idx, do_loads);
+
+#pragma unRoll
+        for (int j = 0; j < T_per_load; j++) {
+            float val = conversion::to<float>(iteration_buffer[j]);
+            val = (val - mean) * denom;
+            val =
+                val * conversion::to<float>(gamma_local[j]) + conversion::to<float>(beta_local[j]);
+            iteration_buffer[j] = conversion::to<T>(val);
+        }
+
+        if (do_loads) {
+            mem_access::store_global<ln::granularity>(block_output + iter_idx, iteration_buffer);
+        }
+    }
+}
+
+#define LAUNCH_FUSED_LN(unRollFactor, threadsPerGroup, maxThreads) \
+    fused_ln<T, unRollFactor, threadsPerGroup, maxThreads>         \
+        <<<grid, block, 0, stream>>>(output, vals, gamma, beta, epsilon, elems_per_row);
+
+template <typename T>
+void launch_fused_ln(T* output,
+                     const T* vals,
+                     const T* gamma,
+                     const T* beta,
+                     float epsilon,
+                     int rows,
+                     int elems_per_row,
+                     cudaStream_t stream)
+{
+    // 8 for __half, 4 for float
+    constexpr int T_per_load = ln::granularity / sizeof(T);
+
+    constexpr int maxThreads = 256;
+
+    // For Flaoat, unRoll 4, for __half, unRoll 2
+    constexpr int internal_unRoll = sizeof(T) == 4 ? 4 : 2;
+
+    const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false;
+    const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internal_unRoll;
+
+    // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of
+    // warp-sized blocks rather than stepping up to 64/96 threads
+    const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step);
+    const int threadsPerGroup = (one_step_threads < maxThreads) ? one_step_threads : maxThreads;
+
+    const int groups_per_block_max =
+        is_subblock_schedule ? (maxThreads + threadsPerGroup - 1) / threadsPerGroup : 1;
+    const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max;
+    const int groups_launch = (groups_per_block + rows - 1) / groups_per_block;
+
+    dim3 block(threadsPerGroup, groups_per_block);
+    dim3 grid(groups_launch);
+
+    const int elems_per_step = threadsPerGroup * h_per_step;
+    const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step;
+
+    if (is_subblock_schedule) {
+        // <=128
+        if (threadsPerGroup == 1) {
+            LAUNCH_FUSED_LN(1, 1, maxThreads);
+        } else if (threadsPerGroup == 2) {
+            LAUNCH_FUSED_LN(1, 2, maxThreads);
+        } else if (threadsPerGroup == 4) {
+            LAUNCH_FUSED_LN(1, 4, maxThreads);
+        } else if (threadsPerGroup == 8) {
+            LAUNCH_FUSED_LN(1, 8, maxThreads);
+        } else if (threadsPerGroup == 16) {
+            LAUNCH_FUSED_LN(1, 16, maxThreads);
+        }
+    } else if (external_unRoll == 1) {
+        // 129 - 4096 elems
+        // (this can launch with 1-7 warps as well)
+        LAUNCH_FUSED_LN(1 * internal_unRoll, maxThreads, maxThreads);
+    } else if (external_unRoll == 2) {
+        // 4097 - 8192 elems
+        LAUNCH_FUSED_LN(2 * internal_unRoll, maxThreads, maxThreads);
+    } else if (external_unRoll == 3) {
+        // 8193 - 12288 elems
+        LAUNCH_FUSED_LN(3 * internal_unRoll, maxThreads, maxThreads);
+    } else if (external_unRoll == 4) {
+        // 12289 - 16384 elems
+        LAUNCH_FUSED_LN(4 * internal_unRoll, maxThreads, maxThreads);
+    }
+}
+
+#define INSTANTIATE_FUSED_LN(T) \
+    template void launch_fused_ln(T*, const T*, const T*, const T*, float, int, int, cudaStream_t);
+
+INSTANTIATE_FUSED_LN(__half);
+#ifdef BF16_AVAILABLE
+INSTANTIATE_FUSED_LN(__nv_bfloat16);
+#endif
+INSTANTIATE_FUSED_LN(float);
+
+/*
+Fused resiual + bias + layer norm implementation. Assumes elems_per_row % 8
+is equal to 0.
+
+TODO(cmikeh2): Goal is to deprecate this implementation. The bias + residual
+need to be fused into compute-bound producer operations.
+
+Args:
+    output: buffer for output data
+    res_output: output of residual addition
+    vals: buffer for input data
+    residual: residual data
+    bias: bias of of input data
+    gamma: gain for normalization
+    beta: bias for normalization
+    epsilon: numeric stability
+    elems_per_row: number of elements each block will normalize
+Template arg:
+    StoreResidual: controls whether the residual calculation is stored
+        or not. When set to false, the input `res_output` is unused.
+*/
+template <typename T, int unRoll, int threadsPerGroup, int maxThreads, bool preLnResidual>
+__global__ void fused_residual_ln(T* output,
+                                  T* res_output,
+                                  const T* vals,
+                                  const T* residual,
+                                  const T* gamma,
+                                  const T* beta,
+                                  float epsilon,
+                                  int elems_per_row)
+{
+    constexpr int T_per_load = ln::granularity / sizeof(T);
+
+    cg::thread_block tb = cg::this_thread_block();
+    cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);
+
+    // X-dimension of the block
+    const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) +
+                             (tb.thread_index().y * elems_per_row);
+    const int thread_offset = tb.thread_index().x * T_per_load;
+    const int base_offset = block_offset + thread_offset;
+    const int stride = tb.size() * T_per_load;
+
+    float sum = reduce::init<rop::Add, float>();
+
+    const T* input_base = vals + base_offset;
+    const T* residual_base = residual + base_offset;
+
+    T local_buffer[unRoll * T_per_load];
+
+    // Unlike a vanilla layernorm, since we're fusing the two adds as well
+    // an inner unRoll seems to be less valuable. If anything, a double unRoll
+    // makes the most sense if we find we are having performance issues.
+#pragma unRoll
+    for (int i = 0; i < unRoll; i++) {
+        T* iteration_buffer = local_buffer + i * T_per_load;
+        T residual_buffer[T_per_load];
+
+        mem_access::load_global<ln::granularity>(
+            iteration_buffer, input_base + i * stride, thread_offset + i * stride < elems_per_row);
+        mem_access::load_global<ln::granularity>(residual_buffer,
+                                                 residual_base + i * stride,
+                                                 thread_offset + i * stride < elems_per_row);
+
+#pragma unRoll
+        for (int j = 0; j < T_per_load; j++) {
+            float vals_up_cast = conversion::to<float>(iteration_buffer[j]);
+            float res_up_cast = conversion::to<float>(residual_buffer[j]);
+            vals_up_cast += res_up_cast;
+            sum = reduce::element<rop::Add>(sum, vals_up_cast);
+            iteration_buffer[j] = conversion::to<T>(vals_up_cast);
+        }
+
+        if (preLnResidual && (thread_offset + i * stride < elems_per_row)) {
+            mem_access::store_global<ln::granularity>(res_output + base_offset + i * stride,
+                                                      iteration_buffer);
+        }
+    }
+
+    reduce::partitioned_block<rop::Add, threadsPerGroup>(tb, warp, sum);
+    const float mean = sum / elems_per_row;
+
+    float mean_diff = reduce::init<rop::Add, float>();
+#pragma unRoll
+    for (int i = 0; i < unRoll; i++) {
+#pragma unRoll
+        for (int j = 0; j < T_per_load; j++) {
+            // Using a 0 value here skews the variance, have to if-guard
+            if (thread_offset + i * stride < elems_per_row) {
+                float diff = (conversion::to<float>(local_buffer[i * T_per_load + j]) - mean);
+                mean_diff = reduce::element<rop::Add>(mean_diff, diff * diff);
+            }
+        }
+    }
+
+    reduce::partitioned_block<rop::Add, threadsPerGroup>(tb, warp, mean_diff);
+    const float variance = mean_diff / elems_per_row;
+    const float denom = __frsqrt_rn(variance + epsilon);
+
+    T* block_output = output + block_offset;
+
+#pragma unRoll
+    for (int i = 0; i < unRoll; i++) {
+        T* iteration_buffer = local_buffer + i * T_per_load;
+        const int iter_idx = i * stride + thread_offset;
+        const bool do_loads = iter_idx < elems_per_row;
+
+        T gamma_local[T_per_load], beta_local[T_per_load];
+
+        mem_access::load_global<ln::granularity>(gamma_local, gamma + iter_idx, do_loads);
+        mem_access::load_global<ln::granularity>(beta_local, beta + iter_idx, do_loads);
+
+#pragma unRoll
+        for (int j = 0; j < T_per_load; j++) {
+            float val = conversion::to<float>(iteration_buffer[j]);
+            val = (val - mean) * denom;
+            val =
+                val * conversion::to<float>(gamma_local[j]) + conversion::to<float>(beta_local[j]);
+            iteration_buffer[j] = conversion::to<T>(val);
+        }
+
+        if (do_loads) {
+            mem_access::store_global<ln::granularity>(block_output + iter_idx, iteration_buffer);
+        }
+    }
+}
+
+// TODO(cmikeh2): There's a bunch of redundancy here that needs to be removed/simplified.
+#define LAUNCH_FUSED_RES_LN(unRollFactor, threadsPerGroup, maxThreads)     \
+    fused_residual_ln<T, unRollFactor, threadsPerGroup, maxThreads, false> \
+        <<<grid, block, 0, stream>>>(                                      \
+            output, nullptr, vals, residual, gamma, beta, epsilon, elems_per_row);
+
+template <typename T>
+void launch_fused_post_ln(T* output,
+                          const T* vals,
+                          const T* residual,
+                          const T* gamma,
+                          const T* beta,
+                          float epsilon,
+                          int rows,
+                          int elems_per_row,
+                          cudaStream_t stream)
+{
+    // 8 for __half, 4 for float
+    constexpr int T_per_load = ln::granularity / sizeof(T);
+
+    constexpr int maxThreads = 256;
+
+    // For Flaoat, unRoll 4, for __half, unRoll 2
+    constexpr int internal_unRoll = sizeof(T) == 4 ? 4 : 2;
+
+    const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false;
+    const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internal_unRoll;
+
+    // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of
+    // warp-sized blocks rather than stepping up to 64/96 threads
+    const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step);
+    const int threadsPerGroup = (one_step_threads < maxThreads) ? one_step_threads : maxThreads;
+
+    const int groups_per_block_max =
+        is_subblock_schedule ? (maxThreads + threadsPerGroup - 1) / threadsPerGroup : 1;
+    const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max;
+    const int groups_launch = (groups_per_block + rows - 1) / groups_per_block;
+
+    dim3 block(threadsPerGroup, groups_per_block);
+    dim3 grid(groups_launch);
+
+    const int elems_per_step = threadsPerGroup * h_per_step;
+    const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step;
+
+    if (is_subblock_schedule) {
+        // <=128
+        if (threadsPerGroup == 1) {
+            LAUNCH_FUSED_RES_LN(1, 1, maxThreads);
+        } else if (threadsPerGroup == 2) {
+            LAUNCH_FUSED_RES_LN(1, 2, maxThreads);
+        } else if (threadsPerGroup == 4) {
+            LAUNCH_FUSED_RES_LN(1, 4, maxThreads);
+        } else if (threadsPerGroup == 8) {
+            LAUNCH_FUSED_RES_LN(1, 8, maxThreads);
+        } else if (threadsPerGroup == 16) {
+            LAUNCH_FUSED_RES_LN(1, 16, maxThreads);
+        }
+    } else if (external_unRoll == 1) {
+        // 129 - 4096 elems
+        // (this can launch with 1-7 warps as well)
+        LAUNCH_FUSED_RES_LN(1 * internal_unRoll, maxThreads, maxThreads);
+    } else if (external_unRoll == 2) {
+        // 4097 - 8192 elems
+        LAUNCH_FUSED_RES_LN(2 * internal_unRoll, maxThreads, maxThreads);
+    } else if (external_unRoll == 3) {
+        // 8193 - 12288 elems
+        LAUNCH_FUSED_RES_LN(3 * internal_unRoll, maxThreads, maxThreads);
+    } else if (external_unRoll == 4) {
+        // 12289 - 16384 elems
+        LAUNCH_FUSED_RES_LN(4 * internal_unRoll, maxThreads, maxThreads);
+    }
+}
+
+#define LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(unRollFactor, threadsPerGroup, maxThreads) \
+    fused_residual_ln<T, unRollFactor, threadsPerGroup, maxThreads, true>               \
+        <<<grid, block, 0, stream>>>(                                                   \
+            norm_output, res_output, vals, residual, gamma, beta, epsilon, elems_per_row);
+
+template <typename T>
+void launch_fused_pre_ln(T* norm_output,
+                         T* res_output,
+                         const T* vals,
+                         const T* residual,
+                         const T* gamma,
+                         const T* beta,
+                         float epsilon,
+                         int rows,
+                         int elems_per_row,
+                         cudaStream_t stream)
+{
+    // 8 for __half, 4 for float
+    constexpr int T_per_load = ln::granularity / sizeof(T);
+
+    constexpr int maxThreads = 256;
+
+    // For Flaoat, unRoll 4, for __half, unRoll 2
+    constexpr int internal_unRoll = sizeof(T) == 4 ? 4 : 2;
+
+    const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false;
+    const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internal_unRoll;
+
+    // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of
+    // warp-sized blocks rather than stepping up to 64/96 threads
+    const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step);
+    const int threadsPerGroup = (one_step_threads < maxThreads) ? one_step_threads : maxThreads;
+
+    const int groups_per_block_max =
+        is_subblock_schedule ? (maxThreads + threadsPerGroup - 1) / threadsPerGroup : 1;
+    const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max;
+    const int groups_launch = (groups_per_block + rows - 1) / groups_per_block;
+
+    dim3 block(threadsPerGroup, groups_per_block);
+    dim3 grid(groups_launch);
+
+    const int elems_per_step = threadsPerGroup * h_per_step;
+    const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step;
+
+    if (is_subblock_schedule) {
+        // <=128
+        if (threadsPerGroup == 1) {
+            LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 1, maxThreads);
+        } else if (threadsPerGroup == 2) {
+            LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 2, maxThreads);
+        } else if (threadsPerGroup == 4) {
+            LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 4, maxThreads);
+        } else if (threadsPerGroup == 8) {
+            LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 8, maxThreads);
+        } else if (threadsPerGroup == 16) {
+            LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1, 16, maxThreads);
+        }
+    } else if (external_unRoll == 1) {
+        // 129 - 4096 elems
+        // (this can launch with 1-7 warps as well)
+        LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(1 * internal_unRoll, maxThreads, maxThreads);
+    } else if (external_unRoll == 2) {
+        // 4097 - 8192 elems
+        LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(2 * internal_unRoll, maxThreads, maxThreads);
+    } else if (external_unRoll == 3) {
+        // 8193 - 12288 elems
+        LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(3 * internal_unRoll, maxThreads, maxThreads);
+    } else if (external_unRoll == 4) {
+        // 12289 - 16384 elems
+        LAUNCH_FUSED_RES_LN_STORE_PRE_LN_RES(4 * internal_unRoll, maxThreads, maxThreads);
+    }
+}
+
+#define INSTANTIATE_RES_LN(T)              \
+    template void launch_fused_post_ln<T>( \
+        T*, const T*, const T*, const T*, const T*, float, int, int, cudaStream_t);
+
+#define INSTANTIATE_PRE_LN_RES(T)         \
+    template void launch_fused_pre_ln<T>( \
+        T*, T*, const T*, const T*, const T*, const T*, float, int, int, cudaStream_t);
+
+INSTANTIATE_RES_LN(__half);
+INSTANTIATE_RES_LN(float);
+#ifdef BF16_AVAILABLE
+INSTANTIATE_RES_LN(__nv_bfloat16);
+#endif
+
+INSTANTIATE_PRE_LN_RES(__half);
+INSTANTIATE_PRE_LN_RES(float);
+#ifdef BF16_AVAILABLE
+INSTANTIATE_PRE_LN_RES(__nv_bfloat16);
+#endif

deepspeed/inference/v2/kernels/core_ops/cuda_linear/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .cuda_linear import *

deepspeed/inference/v2/kernels/core_ops/cuda_linear/cuda_linear.py

@@ -0,0 +1,207 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ....inference_utils import DtypeEnum
+from ....logging import inference_logger
+from deepspeed.ops.op_builder import InferenceCoreBuilder
+from ... import DSKernelBase
+
+
+class CUDAWf6Af16Linear(DSKernelBase):
+    """
+    Wrapper around the CUDA kernel of Wf6Af16 quantized linear.
+
+    Performs z = x @ y
+    """
+    supported_dtypes = [DtypeEnum.fp16]
+
+    def __init__(self):
+        self.inf_module = InferenceCoreBuilder().load()
+        self.inf_module.create_handle()
+        self.kernel = self.inf_module.cuda_wf6af16_linear
+        # The split_k_map is profiled on A100-80G GPU for some common shapes.
+        # It is an array of dictionaries, where the array index is the tokens chunk id.
+        # The dictionary is the mapping from the output channel to the split-K size.
+        self.split_k_map = [
+            {  # tokens: [1, 64]
+                3072: 18,
+                4096: 13,
+                5120: 10,
+                6144: 9,
+                8192: 6,
+                10240: 5,
+                14336: 7,
+                28672: 7,
+                57344: 7
+            },
+            {  # tokens: [65:128]
+                3072: 9,
+                4096: 6,
+                5120: 5,
+                6144: 9,
+                8192: 3,
+                10240: 5,
+                14336: 7,
+                28672: 7,
+                57344: 6
+            },
+            {  # tokens: [129:192]
+                3072: 6,
+                4096: 4,
+                5120: 7,
+                6144: 3,
+                8192: 2,
+                10240: 5,
+                14336: 5,
+                28672: 5,
+                57344: 4
+            },
+            {  # tokens: [193:256]
+                3072: 9,
+                4096: 3,
+                5120: 5,
+                6144: 2,
+                8192: 5,
+                10240: 4,
+                14336: 8,
+                28672: 6,
+                57344: 4
+            },
+            {  # tokens: [257:320]
+                3072: 7,
+                4096: 5,
+                5120: 2,
+                6144: 5,
+                8192: 4,
+                10240: 1,
+                14336: 3,
+                28672: 3,
+                57344: 4
+            },
+            {  # tokens: [321:384]
+                3072: 3,
+                4096: 2,
+                5120: 5,
+                6144: 3,
+                8192: 1,
+                10240: 8,
+                14336: 3,
+                28672: 4,
+                57344: 3
+            },
+            {  # tokens: [385:448]
+                3072: 5,
+                4096: 7,
+                5120: 3,
+                6144: 5,
+                8192: 7,
+                10240: 3,
+                14336: 1,
+                28672: 1,
+                57344: 3
+            },
+            {  # tokens: [449:512]
+                3072: 2,
+                4096: 5,
+                5120: 4,
+                6144: 1,
+                8192: 5,
+                10240: 2,
+                14336: 6,
+                28672: 4,
+                57344: 1
+            },
+            {  # tokens: [513:576]
+                3072: 2,
+                4096: 3,
+                5120: 1,
+                6144: 1,
+                8192: 3,
+                10240: 3,
+                14336: 3,
+                28672: 1,
+                57344: 1
+            },
+            {  # tokens: [577:640]
+                3072: 5,
+                4096: 4,
+                5120: 1,
+                6144: 4,
+                8192: 2,
+                10240: 1,
+                14336: 1,
+                28672: 1,
+                57344: 1
+            },
+            {  # tokens: [641:704]
+                3072: 3,
+                4096: 1,
+                5120: 2,
+                6144: 2,
+                8192: 1,
+                10240: 2,
+                14336: 1,
+                28672: 1,
+                57344: 1
+            },
+            {  # tokens: [705:768]
+                3072: 3,
+                4096: 1,
+                5120: 3,
+                6144: 2,
+                8192: 1,
+                10240: 1,
+                14336: 1,
+                28672: 1,
+                57344: 1
+            }
+        ]
+
+    def __call__(self, output: torch.Tensor, hidden_states: torch.Tensor, weights_2bit: torch.Tensor,
+                 weights_4bit: torch.Tensor, scale: torch.Tensor, out_channels, tokens, in_channels) -> torch.Tensor:
+        """
+        Matmul kernel of FP6 weight-only quantized linear. All inputs should be contiguous.
+        It does not support batched-matmul.
+
+        Parameters:
+            output (torch.Tensor): Output tensor. Shape is of [token_number, out_features]
+            hidden_states (torch.Tensor): Input tensor. Shape is of [token_number, in_features]
+            weights_2bit (torch.Tensor): Input tensor of the 2-bit slice. Shape is of [out_features*2/8, in_features]
+            weights_4bit (torch.Tensor): Input tensor of the 4-bit slice. Shape is of [out_features*4/8, in_features]
+            scale (torch.Tensor): Input tensor. Shape is of [out_features], since the scale is per output channel
+            out_channels (int): The number of output channels
+            tokens (int): The number of tokens
+            in_channels (int): The number of input channels
+        """
+
+        if out_channels % 256 != 0 or in_channels % 64 != 0:
+            raise ValueError("The out and in channel should be multiple of 256 and 64 respectively.")
+
+        # TODO: add a more general heuristic to determine the split-K.
+        split_k = -1  # not initialized
+        if tokens <= 768:
+            # Try to find the split-K from the pre-profiled map.
+            tokens_chunk_id = (tokens - 1) // 64
+            split_k = self.split_k_map[tokens_chunk_id].get(out_channels, -1)
+        if split_k == -1:
+            split_k = 1
+            inference_logger().warning(
+                f"The split-K setting may be suboptimal for shape {tokens}x{in_channels}x{out_channels}...")
+
+        workspace = self.get_workspace(out_channels, tokens, in_channels, split_k, torch.float, hidden_states.device)
+        self.kernel(output, hidden_states, weights_2bit, weights_4bit, scale, workspace, out_channels, tokens,
+                    in_channels, split_k)
+
+    def get_workspace(self, out_channels: int, tokens: int, in_channels: int, split_k: int, dtype,
+                      device) -> torch.Tensor:
+        """
+        Allocate workspace for the kernel. The workspace is used to store the intermediate results of the matmul before
+        split-K. The split-K size is determined by the size of the matmul.
+        """
+        workspace = torch.empty((split_k, out_channels, tokens), dtype=dtype, device=device)
+
+        return workspace

deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/configs.h

@@ -0,0 +1,96 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+// This is a copy of FP6-LLM kernel code: https://arxiv.org/abs/2401.14112
+
+#ifndef CONFIGS_H
+#define CONFIGS_H
+
+// #define DEBUG_MODE
+#define PIPELINE_LEVEL_GMEM 2
+#define PIPELINE_LEVEL_SMEM 2  // only support 2
+
+/************************ Hardware Parameters ************************/
+#define WARP_SIZE 32
+#define REG_BIT_WIDTH 32
+// mma: M=16 K=16 N=8
+#define MMA_8 8
+#define MMA_16 16
+// for memory access
+#define THREAD_OPT_ACCESS_BIT_WIDTH_128 128  // LDS.128, cp_async.128, ...
+#define BIT_WIDTH_PER_HALF 16                // Half precision: FP16
+
+/******************** Register Allocation For GEMM ********************/
+#define REG_PER_THREAD_C_TENSOR_16_16 8  // 8 for FP32 Accumulation
+/********************** Memory Padding Parameters **********************/
+// Eliminating bank-conflict
+#define PADDING_BYTES_16 16  // Padding 16 bytes each column
+#define PADDING_SHARED_MEM_FOR_B_8 \
+    8  // Padding 8 half  each column, during CopyFromGlobalToShared() for B
+#define PADDING_SHARED_MEM_FOR_C_4 \
+    4  // Padding 4 float each column, during StoreToSharedMemoryFromRegister() for C
+/************************* WARP Tiling part-1 *************************/
+#define WARP_ROW_MMA_TENSORS 4
+#define WARP_M (WARP_ROW_MMA_TENSORS * MMA_16)  // 64
+#define WARP_K_MMA_TENSORS 4
+#define WARP_K (WARP_K_MMA_TENSORS * MMA_16)  // 64
+template <int BLOCK_ROW_WARPS_, int BLOCK_COL_WARPS_, int WARP_COL_MMA_TENSORS_>
+struct TilingConfig {
+    // Depending on "n" dimension of the GEMM
+    static constexpr int BLOCK_ROW_WARPS = BLOCK_ROW_WARPS_;
+    static constexpr int BLOCK_COL_WARPS = BLOCK_COL_WARPS_;
+    static constexpr int WARP_COL_MMA_TENSORS = WARP_COL_MMA_TENSORS_;
+    /************************* WARP Tiling part-2 *************************/
+    static constexpr int WARP_N = WARP_COL_MMA_TENSORS * MMA_8;
+    /*************************Thread Block Tiling *************************/
+    static constexpr int TILE_M = WARP_M * BLOCK_ROW_WARPS;
+    static constexpr int TILE_N = MMA_8 * WARP_COL_MMA_TENSORS * BLOCK_COL_WARPS;
+    static constexpr int TILE_K = WARP_K;
+    /********************** #Thread per Thread Block **********************/
+    static constexpr int BLOCK_WARPS = BLOCK_ROW_WARPS * BLOCK_COL_WARPS;
+    static constexpr int BLOCK_THREADS = BLOCK_WARPS * WARP_SIZE;
+    /******************************* Others *******************************/
+    static constexpr int SMEM_SIZE_B_TILE = TILE_N * (TILE_K + PADDING_BYTES_16) * 2 *
+                                            PIPELINE_LEVEL_GMEM;  // sizeof(half)=2, doubleBuffer=2
+    static constexpr int SMEM_SIZE_C_TILE =
+        TILE_N * (TILE_M + PADDING_BYTES_16) * 4;  // sizeof(float)=4
+};
+
+/************************ General Config for Quant-LLM **********************/
+#define WEIGHT_FRAG1_BIT_WIDTH 2
+#define WEIGHT_FRAG2_BIT_WIDTH 4
+#define WEIGHT_BIT_WIDTH (WEIGHT_FRAG1_BIT_WIDTH + WEIGHT_FRAG2_BIT_WIDTH)  // 6
+// #define QUANT_GROUP_SIZE_DIVIDED_BY_64  4                                                   //
+//  QuantGroupSize: 4*64 = 256
+/*************************** 64*64 Weghts of A WARP *************************/
+#define WEIGHT_PER_UNIT (WARP_M * WARP_K)  // 64*64
+#define SMEM_SIZE_IN_BYTES_PER_WARP_A1          \
+    (WEIGHT_PER_UNIT * WEIGHT_FRAG1_BIT_WIDTH / \
+     8)  // 1024 Bytes   #doubleBuffer not takedn into consideration
+#define SMEM_SIZE_IN_BYTES_PER_WARP_A2          \
+    (WEIGHT_PER_UNIT * WEIGHT_FRAG2_BIT_WIDTH / \
+     8)  // 2048 Bytes   #doubleBuffer not takedn into consideration
+#define SMEM_SIZE_A1_TILE                 \
+    (SMEM_SIZE_IN_BYTES_PER_WARP_A1 * 4 * \
+     PIPELINE_LEVEL_GMEM)  // #WARP=4, #Trible-Buffer for 3-level pipeline for A = 12 KB; double
+                           // buffer for 2-level pipeline A= 8  KB.
+#define SMEM_SIZE_A2_TILE                 \
+    (SMEM_SIZE_IN_BYTES_PER_WARP_A2 * 4 * \
+     PIPELINE_LEVEL_GMEM)  // #WARP=4, #Trible-Buffer for 3-level pipeline for A = 24 KB; double
+                           // buffer for 2-level pipeline A= 16 KB.
+/******************** Global Memory Layout For QUANTIZED DATA ******************/
+#define NUM_INT4_PER_UNIT_2BIT_FRAG (WEIGHT_PER_UNIT * WEIGHT_FRAG1_BIT_WIDTH / 128)  // 64
+#define NUM_INT4_PER_UNIT_4BIT_FRAG (WEIGHT_PER_UNIT * WEIGHT_FRAG2_BIT_WIDTH / 128)  // 128
+/******************** Register Allocation For QUANTIZED DATA ******************/
+#define WEIGHT_PER_THREAD (WEIGHT_PER_UNIT / WARP_SIZE)                   // 128
+#define REG_PER_THREAD_2BIT_FRAG (WEIGHT_PER_THREAD / REG_BIT_WIDTH * 2)  // 8
+#define REG_PER_THREAD_4BIT_FRAG (WEIGHT_PER_THREAD / REG_BIT_WIDTH * 4)  // 16
+/******************** Register Allocation For QUANT Scales ******************/
+#define WARP_REG_QUANT_SCALE 4  // 8 rows per thread -> 8 FP16 scales -> 4 registers
+#define WARP_REG_QUANT_SCALE_DISTRIBUTED \
+    1  // T0-T3, T4-T7, ..., T28-T31 share the same scales, using shfl to get all the scales for
+       // each thread
+
+#endif  // CONFIGS_H

deepspeed/inference/v2/kernels/core_ops/cuda_linear/include/weight_prepacking.h

@@ -0,0 +1,207 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+// This is a copy of FP6-LLM kernel code: https://arxiv.org/abs/2401.14112
+
+#ifndef DEEPSPEED_CUDA_LINEAR_WEIGHT_PREPACKING_H
+#define DEEPSPEED_CUDA_LINEAR_WEIGHT_PREPACKING_H
+
+#include <assert.h>
+#include <stdio.h>
+#include <vector>
+
+void Padding_8_FP6_To_8_Bytes(unsigned char Padded_FP6[],
+                              unsigned char* FP6_Array)  // padding 0 to the lowerest bit location
+{
+    Padded_FP6[0] = FP6_Array[0] & 0xfc;
+    Padded_FP6[1] = (FP6_Array[0] << 6) | ((FP6_Array[1] >> 2) & 0xfc);
+    Padded_FP6[2] = (FP6_Array[1] << 4) | ((FP6_Array[2] >> 4) & 0xfc);
+    Padded_FP6[3] = FP6_Array[2] << 2;
+    Padded_FP6[4] = FP6_Array[3] & 0xfc;
+    Padded_FP6[5] = (FP6_Array[3] << 6) | ((FP6_Array[4] >> 2) & 0xfc);
+    Padded_FP6[6] = (FP6_Array[4] << 4) | ((FP6_Array[5] >> 4) & 0xfc);
+    Padded_FP6[7] = FP6_Array[5] << 2;
+}
+
+unsigned char Extract_2_Bits_From_4_PaddedFP6(unsigned char B1,
+                                              unsigned char B2,
+                                              unsigned char B3,
+                                              unsigned char B4)
+{
+    unsigned char out;
+    out = (B1 & 0xc0) | ((B2 & 0xc0) >> 2) | ((B3 & 0xc0) >> 4) | ((B4 & 0xc0) >> 6);
+    return out;
+}
+
+unsigned char Extract_4_Bits_From_2_PaddedFP6(
+    unsigned char B1,
+    unsigned char
+        B2)  // The highest two bits are already extracted by Extract_2_Bits_From_4_PaddedFP6();
+{
+    unsigned char out;
+    out = ((B1 << 2) & 0xf0) | ((B2 >> 2) & 0x0f);
+    return out;
+}
+
+// dealing with 4 1*8 blocks of FP6
+void Assign_32_FP6_To_4_Thread(std::vector<unsigned char> Seg_2bit[],
+                               std::vector<unsigned char> Seg_4bit[],
+                               unsigned char* PTR_1,
+                               unsigned char* PTR_2,
+                               unsigned char* PTR_3,
+                               unsigned char* PTR_4)
+{
+    unsigned char Padded_8_FP8[4][8];
+    Padding_8_FP6_To_8_Bytes(Padded_8_FP8[0], PTR_1);
+    Padding_8_FP6_To_8_Bytes(Padded_8_FP8[1], PTR_2);
+    Padding_8_FP6_To_8_Bytes(Padded_8_FP8[2], PTR_3);
+    Padding_8_FP6_To_8_Bytes(Padded_8_FP8[3], PTR_4);
+    //
+    unsigned char Seg1_Byte1_T[4];
+    unsigned char Seg1_Byte2_T[4];
+    unsigned char Seg2_Byte1_T[4];
+    unsigned char Seg2_Byte2_T[4];
+    unsigned char Seg2_Byte3_T[4];
+    unsigned char Seg2_Byte4_T[4];
+    for (int t = 0; t < 4; t++) {
+        Seg1_Byte1_T[t] = Extract_2_Bits_From_4_PaddedFP6(Padded_8_FP8[0][0 + t * 2],
+                                                          Padded_8_FP8[0][1 + t * 2],
+                                                          Padded_8_FP8[1][0 + t * 2],
+                                                          Padded_8_FP8[1][1 + t * 2]);
+        Seg1_Byte2_T[t] = Extract_2_Bits_From_4_PaddedFP6(Padded_8_FP8[2][0 + t * 2],
+                                                          Padded_8_FP8[2][1 + t * 2],
+                                                          Padded_8_FP8[3][0 + t * 2],
+                                                          Padded_8_FP8[3][1 + t * 2]);
+        Seg2_Byte1_T[t] =
+            Extract_4_Bits_From_2_PaddedFP6(Padded_8_FP8[0][0 + t * 2], Padded_8_FP8[0][1 + t * 2]);
+        Seg2_Byte2_T[t] =
+            Extract_4_Bits_From_2_PaddedFP6(Padded_8_FP8[1][0 + t * 2], Padded_8_FP8[1][1 + t * 2]);
+        Seg2_Byte3_T[t] =
+            Extract_4_Bits_From_2_PaddedFP6(Padded_8_FP8[2][0 + t * 2], Padded_8_FP8[2][1 + t * 2]);
+        Seg2_Byte4_T[t] =
+            Extract_4_Bits_From_2_PaddedFP6(Padded_8_FP8[3][0 + t * 2], Padded_8_FP8[3][1 + t * 2]);
+    }
+    //
+    for (int t = 0; t < 4; t++) {
+        Seg_2bit[t].push_back(Seg1_Byte1_T[t]);
+        Seg_2bit[t].push_back(Seg1_Byte2_T[t]);
+        Seg_4bit[t].push_back(Seg2_Byte1_T[t]);
+        Seg_4bit[t].push_back(Seg2_Byte2_T[t]);
+        Seg_4bit[t].push_back(Seg2_Byte3_T[t]);
+        Seg_4bit[t].push_back(Seg2_Byte4_T[t]);
+    }
+    return;
+}
+
+void BitInterleaving_2bit(unsigned char* PTR_4Bytes)
+{
+    unsigned int* PTR_UINT = reinterpret_cast<unsigned int*>(PTR_4Bytes);
+    unsigned int input = *PTR_UINT;
+    //
+    // int order_2bit[16] = {1,5,9,13,3,7,11,15,2,6,10,14,4,8,12,16};  // pre-defined order for
+    // bit-interleaving in QuantLLM
+    int order_2bit[16] = {
+        2, 6, 10, 14, 4, 8, 12, 16, 1, 5, 9, 13, 3, 7, 11, 15};  // pre-defined order for
+                                                                 // bit-interleaving in QuantLLM
+    unsigned int Frags_2bit[16];  // The highest 2 bits are used to store the extracted fragments.
+    for (int i = 0; i < 16; i++) Frags_2bit[i] = (input << 2 * (order_2bit[i] - 1)) & 0xc0000000;
+    //
+    unsigned int output = 0x00000000;
+    for (int i = 0; i < 16; i++) output |= (Frags_2bit[i] >> (i * 2));
+    //
+    *PTR_UINT = output;
+}
+
+void BitInterleaving_4bit(unsigned char* PTR_4Bytes)
+{
+    unsigned int* PTR_UINT = reinterpret_cast<unsigned int*>(PTR_4Bytes);
+    unsigned int input = *PTR_UINT;
+    //
+    // int order_4bit[8] = {1,5,3,7,2,6,4,8};  // pre-defined order for bit-interleaving in QuantLLM
+    int order_4bit[8] = {
+        2, 6, 4, 8, 1, 5, 3, 7};  // pre-defined order for bit-interleaving in QuantLLM
+    unsigned int Frags_4bit[8];   // The highest4 bits are used to store the extracted fragments.
+    for (int i = 0; i < 8; i++) Frags_4bit[i] = (input << 4 * (order_4bit[i] - 1)) & 0xf0000000;
+    //
+    unsigned int output = 0x00000000;
+    for (int i = 0; i < 8; i++) output |= (Frags_4bit[i] >> (i * 4));
+    //
+    *PTR_UINT = output;
+}
+
+/*
+ * Inputs:
+ * (1) unsigned char Weight_6bit [M*K*6/8]
+ * Outputs:
+ * (1) unsigned char Weight_2bit [M*K*2/8]
+ * (2) unsigned char Weight_4bit [M*K*4/8]
+ *
+ * Assumption: Weight_6bit, Weight_2bit, Weight_4bit all stored continuously in row-major.
+ * 8 FP6 = 6 Bytes
+ * 8 FP4 = 4 Bytes
+ * 8 FP2 = 2 Bytes
+ */
+void weight_matrix_prepacking(int* FP6Weights, size_t M, size_t K)
+{
+    assert(M % 64 == 0);
+    assert(K % 64 == 0);
+    //
+    unsigned char* Weight_6bit = reinterpret_cast<unsigned char*>(FP6Weights);
+    unsigned char* Weight_2bit = Weight_6bit;
+    unsigned char* Weight_4bit = Weight_6bit + M * K * 2 / 8;
+    //
+    std::vector<unsigned char> A_Segment_2bit[32];
+    std::vector<unsigned char> A_Segment_4bit[32];
+    //
+    size_t BytesPerRow = K * 6 / 8;
+    // Pass-1: (1) 2+4 split; (2) assign weights to 32 threads.
+    for (size_t i = 0; i < M / 64; i++)  //
+    {
+        for (size_t j = 0; j < K / 16; j++) {
+            for (size_t k = 0; k < 64 / 16; k++) {
+                size_t row = i * 64 + k * 16;
+                size_t col = j * 16;
+                unsigned char* StartPTR_1 = Weight_6bit + row * BytesPerRow + col * 6 / 8;
+                unsigned char* StartPTR_2 = StartPTR_1 + 8 * BytesPerRow;
+                unsigned char* StartPTR_3 = StartPTR_1 + 8 * 6 / 8;
+                unsigned char* StartPTR_4 = StartPTR_2 + 8 * 6 / 8;
+                // Dealing with each 16*16 blocks then...
+                for (int l = 0; l < 8; l++)
+                    Assign_32_FP6_To_4_Thread(&A_Segment_2bit[l * 4],
+                                              &A_Segment_4bit[l * 4],
+                                              StartPTR_1 + l * BytesPerRow,
+                                              StartPTR_2 + l * BytesPerRow,
+                                              StartPTR_3 + l * BytesPerRow,
+                                              StartPTR_4 + l * BytesPerRow);
+            }
+        }
+    }
+    // Verifying the length of 2_bit segments and 4_bit segments
+    size_t BytesPerThread_2bit = M * K * 2 / 8 / 32;
+    size_t BytesPerThread_4bit = M * K * 4 / 8 / 32;
+    for (int i = 0; i < 32; i++) {
+        assert(A_Segment_2bit[i].size() == BytesPerThread_2bit);
+        assert(A_Segment_4bit[i].size() == BytesPerThread_4bit);
+    }
+    // Pass-2: Optimizing coleasced global memory access
+    for (size_t i = 0; i < BytesPerThread_2bit / 4; i++)
+        for (int t = 0; t < 32; t++)
+            for (int b = 0; b < 4; b++)
+                Weight_2bit[i * 128 + t * 4 + (3 - b)] =
+                    A_Segment_2bit[t]
+                                  [i * 4 + b];  // why (3-b): special byte order within a register
+    for (size_t i = 0; i < BytesPerThread_4bit / 4; i++)
+        for (int t = 0; t < 32; t++)
+            for (int b = 0; b < 4; b++)
+                Weight_4bit[i * 128 + t * 4 + (3 - b)] =
+                    A_Segment_4bit[t][i * 4 + b];  // why (3-b):special byte order within a register
+    // Pass-3: Bit-level interleaving
+    for (size_t i = 0; i < BytesPerThread_2bit * 32 / 4; i++)
+        BitInterleaving_2bit(Weight_2bit + 4 * i);
+    for (size_t i = 0; i < BytesPerThread_4bit * 32 / 4; i++)
+        BitInterleaving_4bit(Weight_4bit + 4 * i);
+}
+
+#endif

deepspeed/inference/v2/kernels/core_ops/cuda_linear/linear_kernels.cpp

@@ -0,0 +1,224 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include <ATen/cuda/CUDAContext.h>
+
+#include "linear_kernels.h"
+
+namespace {
+
+// For bit-level debugging.
+template <typename T>
+void print_bits(T num)
+{
+    char bits[sizeof(T) * 8 + 1] = {'\0'};
+    for (int bit = 0; bit < (sizeof(T) * 8); bit++) {
+        bits[sizeof(T) * 8 - 1 - bit] = '0' + (num & 0x01);
+        num = num >> 1;
+    }
+    printf("%s\n", bits);
+}
+
+void print_bits(half num)
+{
+    char bits[sizeof(half) * 8 + 1] = {'\0'};
+    auto int_num = *reinterpret_cast<uint16_t*>(&num);
+    for (int bit = 0; bit < (sizeof(half) * 8); bit++) {
+        bits[sizeof(half) * 8 - 1 - bit] = '0' + (int_num & 0x01);
+        int_num = int_num >> 1;
+    }
+    printf("%s\n", bits);
+}
+
+/*
+ * Function to pack 4 fake quantized FP16 value into continuously stored 4 FP6 values.
+ */
+void cast_fp16_fp6(uint16_t* FP16x4, uint8_t* FP6x4)
+{
+    // Constants for FP6
+    constexpr int exponent_nbits_fp6 = 3;
+    constexpr int mantissa_nbits_fp6 = 2;
+    constexpr int exp_bias_fp6 = (1 << (exponent_nbits_fp6 - 1)) - 1;
+    // Constants for FP16
+    constexpr int exponent_nbits_fp16 = 5;
+    constexpr int mantissa_nbits_fp16 = 10;
+    constexpr int exp_bias_fp16 = (1 << (exponent_nbits_fp16 - 1)) - 1;
+
+    int fp6_temp[4];
+
+    float absmin_nonzero_fp6 = 0.0625;
+    // Note that we regard the exponent of '111' as a regular value rather than NaN or inf. This is
+    // the same with that in qtorch.
+    float absmax_fp6 = 28;
+
+    for (int i = 0; i < 4; ++i) {
+        uint16_t source = FP16x4[i];
+        float fp6_value_abs = std::abs(__half2float(*((half*)(&source))));
+        if ((fp6_value_abs != 0 && fp6_value_abs < absmin_nonzero_fp6) ||
+            fp6_value_abs > absmax_fp6) {
+            // TODO(zhen): a better way may be rounding it to the nearest FP6 value.
+            throw std::invalid_argument("Input value out of range for FP6.");
+        }
+
+        // It is not safe to do shift operation on uint16_t. So we promote it to int.
+        int source_promote = int(source);
+
+        int sign_bit = (source_promote >> 15);
+        // Extracting exponent represented in FP16. The sign mask 0x7FFF is '0111 1111 1111 1111'
+        int exp_bit = (source_promote & 0x7FFF) >> mantissa_nbits_fp16;
+        // Extracting mantissa represented in FP16
+        int mant_bit = source_promote & ((1 << mantissa_nbits_fp16) - 1);
+
+        int new_exp_bit;
+        int new_mant_bit;
+
+        if (exp_bit == 0) {
+            // Subnormal FP16 number. Too small for FP6.
+            new_exp_bit = 0;
+            new_mant_bit = 0;
+        } else {
+            new_mant_bit = mant_bit >> (mantissa_nbits_fp16 - mantissa_nbits_fp6);
+            new_exp_bit = exp_bit - exp_bias_fp16 + exp_bias_fp6;
+
+            // Deal with subnormal FP6 values.
+            int target_exp_val = exp_bit - exp_bias_fp16;
+            int min_fp6_exp_val = -exp_bias_fp6 + 1;
+            bool subnormal_fp6 = target_exp_val < min_fp6_exp_val;
+            if (subnormal_fp6) {
+                // TODO(zhen): add the rounding logic.
+                new_exp_bit = 0;
+                // The implicit 1 in the mantissa of FP16 is not present in subnormal FP6. Thus we
+                // need to add it
+                new_mant_bit = (new_mant_bit | (1 << mantissa_nbits_fp6)) >>
+                               (min_fp6_exp_val - target_exp_val);
+            }
+        }
+
+        fp6_temp[i] = (sign_bit << (exponent_nbits_fp6 + mantissa_nbits_fp6)) |
+                      (new_exp_bit << mantissa_nbits_fp6) | new_mant_bit;
+    }
+    // Pack the values
+    FP6x4[0] = fp6_temp[0] << 2 | (fp6_temp[1] >> 4);
+    FP6x4[1] = (fp6_temp[1] & 0x0F) << 4 | (fp6_temp[2] >> 2);
+    FP6x4[2] = (fp6_temp[2] & 0x03) << 6 | fp6_temp[3];
+}
+
+/*
+ *  Function to prepack FP16 weights into continuous FP6 values.
+ *
+ *  Parameters:
+ *     weight_16bit: input weight in FP16, size M*K
+ *     weight_6bit: output weight in packed FP6, continuously stored, size M*K*6/8
+ *     M, K: the shape of the weight
+ */
+void weight_prepacking_fp16_to_fp6(uint16_t* weight_16bit,
+                                   uint8_t* weight_6bit_packed,
+                                   size_t M,
+                                   size_t K)
+{
+    // Every four 16-bit elements are packed into three 6-bit values (4*6bit == 3*8bit).
+    if (K * 6 % 8 != 0) { throw std::invalid_argument("(K * 6 % 8) should be 0"); }
+    size_t K_fp6_packed = K * 6 / 8;
+    // #pragma omp parallel for
+    for (auto m = 0; m < M; m++) {
+        uint8_t* ptr_6bit = weight_6bit_packed + m * K_fp6_packed;
+        uint16_t* ptr_16bit = weight_16bit + m * K;
+        for (auto k = 0; k < K; k += 4) {
+            cast_fp16_fp6(ptr_16bit, ptr_6bit);
+            ptr_16bit += 4;
+            ptr_6bit += 3;
+        }
+    }
+}
+
+}  // namespace
+
+/*
+ * Function to execute the FP6 linear kernel.
+ *
+ * Parameters:
+ *    output: output tensor, size M*N
+ *    hidden_states: input activation tensor, size N*K
+ *    weights_2bit: packed 2bit weights, size M*K*2/8
+ *    weights_4bit: packed 4bit weights, size M*K*4/8
+ *    scales: scale tensor, size M
+ *    workspace: workspace tensor, size M*N*split_k
+ *    M: the output channel number of the weight
+ *    N: the token number of the activation
+ *    K: the input channel number of the weight
+ *    split_k: the split size of the GEMM calculation
+ */
+void cuda_wf6af16_linear(torch::Tensor& output,
+                         torch::Tensor& hidden_states,
+                         torch::Tensor& weights_2bit,
+                         torch::Tensor& weights_4bit,
+                         torch::Tensor& scales,
+                         torch::Tensor& workspace,
+                         int M,
+                         int N,
+                         int K,
+                         int split_k)
+{
+    TORCH_CHECK(weights_2bit.device().type() == torch::kCUDA, "weight_2bit must be on CUDA");
+    TORCH_CHECK(weights_4bit.device().type() == torch::kCUDA, "weight_4bit must be on CUDA");
+    TORCH_CHECK(hidden_states.device().type() == torch::kCUDA, "X must be on CUDA");
+    TORCH_CHECK(scales.device().type() == torch::kCUDA, "scales must be on CUDA");
+
+    auto status = fp6_linear_kernel(at::cuda::getCurrentCUDAStream(),
+                                    (uint4*)(weights_2bit.data_ptr<uint8_t>()),
+                                    (uint4*)(weights_4bit.data_ptr<uint8_t>()),
+                                    (half*)(scales.data_ptr<at::Half>()),
+                                    (half*)(hidden_states.data_ptr<at::Half>()),
+                                    (half*)(output.data_ptr<at::Half>()),
+                                    M,
+                                    N,
+                                    K,
+                                    workspace.data_ptr<float>(),
+                                    split_k);
+    if (status != cudaSuccess) {
+        AT_ERROR("fp6_linear_kernel failed with error: ", cudaGetErrorString(status));
+    }
+}
+
+/*
+ * Function to prepack the fake 6-bit-quantized FP16 weights into 2bit and 4bit.
+ *
+ * Parameters:
+ *    weight: input weight in FP16 (containing the quantized FP6-ranged value), size M*K
+ * Returns:
+ *   weight_2bit: output weight in 2bit, size M*K*2/8
+ *   weight_4bit: output weight in 4bit, size M*K*4/8
+ */
+std::vector<torch::Tensor> preprocess_weight(torch::Tensor& weight)
+{
+    TORCH_CHECK(weight.dim() == 2, "weight must be 2-dimensional");
+    TORCH_CHECK(weight.scalar_type() == torch::kFloat16, "weight must be FP16");
+    TORCH_CHECK(weight.is_contiguous(), "weight must be contiguous");
+    TORCH_CHECK(weight.device().type() == torch::kCPU, "weight must be on CPU");
+    auto M = weight.size(0);
+    auto K = weight.size(1);
+    TORCH_CHECK(K % 4 == 0, "K must be multiple of 4");
+
+    // Pack weight from FP16 to FP6.
+    uint16_t* weight_16bit_ptr = reinterpret_cast<uint16_t*>(weight.data_ptr<at::Half>());
+    std::vector<uint8_t> weight_6bit_packed(M * K * 6 / 8);
+    uint8_t* weight_6bit_ptr = weight_6bit_packed.data();
+    weight_prepacking_fp16_to_fp6(weight_16bit_ptr, weight_6bit_ptr, M, K);
+
+    // Split weight into 2bit and 4bit.
+    weight_matrix_prepacking(reinterpret_cast<int*>(weight_6bit_ptr), M, K);
+    uint8_t* weight_2bit_ptr = weight_6bit_ptr;
+
+    // Make sure that the new split tensor does not share the underlying memory with the original
+    // one. Otherwise it will incur some problems when the original tensor is deleted. It also
+    // makes the memory flattern risky.
+    auto weight_2bit =
+        torch::from_blob(weight_2bit_ptr, {M * K * 2 / 8}, torch::kUInt8).clone().detach();
+    uint8_t* weight_4bit_ptr = weight_2bit_ptr + M * K * 2 / 8;
+    auto weight_4bit =
+        torch::from_blob(weight_4bit_ptr, {M * K * 4 / 8}, torch::kUInt8).clone().detach();
+
+    return {weight_2bit, weight_4bit};
+}

deepspeed/inference/v2/kernels/core_ops/cuda_linear/linear_kernels.h

@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#ifndef DEEPSPEED_CUDA_LINEAR_KERNELS_H
+#define DEEPSPEED_CUDA_LINEAR_KERNELS_H
+
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+#include "ds_kernel_utils.h"
+
+#include "linear_kernels_cuda.h"
+
+void cuda_wf6af16_linear(torch::Tensor& output,
+                         torch::Tensor& hidden_states,
+                         torch::Tensor& weights_2bit,
+                         torch::Tensor& weights_4bit,
+                         torch::Tensor& scale,
+                         torch::Tensor& workspace,
+                         int M,
+                         int N,
+                         int K,
+                         int split_k);
+
+std::vector<torch::Tensor> preprocess_weight(torch::Tensor& Weight);
+
+#endif

deepspeed/inference/v2/kernels/core_ops/cuda_linear/linear_kernels_cuda.cu

@@ -0,0 +1,318 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+// This is a copy of FP6-LLM kernel code: https://arxiv.org/abs/2401.14112
+
+// clang-format off
+// Put the torch headers at the front to avoid conflict with other headers on
+// `at::nullopt` and `at::optional`.
+#include <torch/extension.h>
+#include <ATen/ATen.h>
+// clang-format on
+
+#include "include/kernel_matmul.cuh"
+#include "include/kernel_reduction.cuh"
+#include "include/weight_prepacking.h"
+
+#include <assert.h>
+#include <stdio.h>
+
+#include "linear_kernels_cuda.h"
+
+template <typename TilingConfig, typename OutputDataType>
+static void Kernel_Ex(cudaStream_t stream,
+                      const uint4* Weight1,
+                      const uint4* Weight2,
+                      const half* Scales,
+                      const half* B,
+                      OutputDataType* C,
+                      const size_t M_Global,
+                      const size_t N_Global,
+                      const size_t K_Global,
+                      int Split_K)
+{
+#ifdef DEBUG_MODE
+    printf("\n");
+    printf("Launcher.cu->Kernel_Ex():\n");
+    printf("M: %d, N: %d, K: %d, SplitK: %d\n", M_Global, N_Global, K_Global, Split_K);
+    printf("TILE_M: %d, TILE_K: %d, TILE_N: %d\n",
+           TilingConfig::TILE_M,
+           TilingConfig::TILE_K,
+           TilingConfig::TILE_N);
+#endif
+    static size_t SHMEM_SZ =
+        max(TilingConfig::SMEM_SIZE_B_TILE + SMEM_SIZE_A1_TILE + SMEM_SIZE_A2_TILE,
+            TilingConfig::SMEM_SIZE_C_TILE);
+    cudaFuncSetAttribute(QUANT_GEMM_Kernel<TilingConfig, OutputDataType>,
+                         cudaFuncAttributeMaxDynamicSharedMemorySize,
+                         SHMEM_SZ);
+    size_t dimN = (N_Global - 1) / TilingConfig::TILE_N + 1;
+    size_t dimM = M_Global * Split_K / TilingConfig::TILE_M;
+    dim3 GridDim(dimN, dimM, 1);
+    dim3 BlockDim(WARP_SIZE * TilingConfig::BLOCK_WARPS, 1, 1);
+
+#ifdef DEBUG_MODE
+    printf(
+        "GridDim.x: %d, GridDim.y: %d, GridDim.z: %d, BlockDim.x: %d, BlockDim.y: %d, BlockDim.z: "
+        "%d SHMEM_SZ: %d\n",
+        GridDim.x,
+        GridDim.y,
+        GridDim.z,
+        BlockDim.x,
+        BlockDim.y,
+        BlockDim.z,
+        SHMEM_SZ);
+    printf("\n");
+#endif
+
+    QUANT_GEMM_Kernel<TilingConfig, OutputDataType><<<GridDim, BlockDim, SHMEM_SZ, stream>>>(
+        Weight1, Weight2, Scales, B, C, M_Global, N_Global, K_Global, Split_K);
+}
+
+/*
+ *
+ */
+cudaError_t fp6_linear_kernel(cudaStream_t stream,
+                              const uint4* Weight1,
+                              const uint4* Weight2,
+                              const half* Scales,
+                              const half* B,
+                              half* C,
+                              const size_t M_Global,
+                              const size_t N_Global,
+                              const size_t K_Global,
+                              float* Reduction_Workspace,  // Reduction_Workspace_Size = Split_K *
+                                                           // M_Global * N_Global * sizeof(fp32)
+                              int Split_K)
+{
+    assert(M_Global % 256 == 0);
+    assert(K_Global % 64 == 0);
+    assert(N_Global > 0);
+
+    // Work around to support more N shapes:
+    size_t N_PowerOf2;
+    if (N_Global > 0 && N_Global <= 8) N_PowerOf2 = 8;
+    if (N_Global > 8 && N_Global <= 16) N_PowerOf2 = 16;
+    if (N_Global > 16 && N_Global <= 32) N_PowerOf2 = 32;
+    if (N_Global > 32 && N_Global <= 64) N_PowerOf2 = 64;
+    if (N_Global > 64 && N_Global <= 128) N_PowerOf2 = 128;
+    if (N_Global > 128) N_PowerOf2 = ((N_Global - 1) / 128 + 1) * 128;
+
+    if (Split_K == 1) {
+        switch (N_PowerOf2) {
+            case 8:
+                Kernel_Ex<TilingConfig<4, 1, 1>, half>(
+                    stream, Weight1, Weight2, Scales, B, C, M_Global, N_Global, K_Global, Split_K);
+                break;
+            case 16:
+                Kernel_Ex<TilingConfig<4, 1, 2>, half>(
+                    stream, Weight1, Weight2, Scales, B, C, M_Global, N_Global, K_Global, Split_K);
+                break;
+            case 32:
+                Kernel_Ex<TilingConfig<4, 1, 4>, half>(
+                    stream, Weight1, Weight2, Scales, B, C, M_Global, N_Global, K_Global, Split_K);
+                break;
+            case 64:
+                Kernel_Ex<TilingConfig<4, 1, 8>, half>(
+                    stream, Weight1, Weight2, Scales, B, C, M_Global, N_Global, K_Global, Split_K);
+                break;
+            case 128:
+                Kernel_Ex<TilingConfig<4, 1, 8>, half>(
+                    stream, Weight1, Weight2, Scales, B, C, M_Global, N_Global, K_Global, Split_K);
+                break;
+            default:
+                if (N_PowerOf2 % 128 != 0) {
+                    printf("QuantLLM_API Error: Unsupported N dimension %lu!\n", N_PowerOf2);
+                    return cudaErrorUnknown;
+                }
+                Kernel_Ex<TilingConfig<4, 1, 8>, half>(
+                    stream, Weight1, Weight2, Scales, B, C, M_Global, N_Global, K_Global, Split_K);
+                break;
+        }
+    } else {
+        switch (N_PowerOf2) {
+            case 8:
+                Kernel_Ex<TilingConfig<4, 1, 1>, float>(stream,
+                                                        Weight1,
+                                                        Weight2,
+                                                        Scales,
+                                                        B,
+                                                        Reduction_Workspace,
+                                                        M_Global,
+                                                        N_Global,
+                                                        K_Global,
+                                                        Split_K);
+                break;
+            case 16:
+                Kernel_Ex<TilingConfig<4, 1, 2>, float>(stream,
+                                                        Weight1,
+                                                        Weight2,
+                                                        Scales,
+                                                        B,
+                                                        Reduction_Workspace,
+                                                        M_Global,
+                                                        N_Global,
+                                                        K_Global,
+                                                        Split_K);
+                break;
+            case 32:
+                Kernel_Ex<TilingConfig<4, 1, 4>, float>(stream,
+                                                        Weight1,
+                                                        Weight2,
+                                                        Scales,
+                                                        B,
+                                                        Reduction_Workspace,
+                                                        M_Global,
+                                                        N_Global,
+                                                        K_Global,
+                                                        Split_K);
+                break;
+            case 64:
+                Kernel_Ex<TilingConfig<4, 1, 8>, float>(stream,
+                                                        Weight1,
+                                                        Weight2,
+                                                        Scales,
+                                                        B,
+                                                        Reduction_Workspace,
+                                                        M_Global,
+                                                        N_Global,
+                                                        K_Global,
+                                                        Split_K);
+                break;
+            case 128:
+                Kernel_Ex<TilingConfig<4, 1, 8>, float>(stream,
+                                                        Weight1,
+                                                        Weight2,
+                                                        Scales,
+                                                        B,
+                                                        Reduction_Workspace,
+                                                        M_Global,
+                                                        N_Global,
+                                                        K_Global,
+                                                        Split_K);
+                break;
+            default:
+                if (N_PowerOf2 % 128 != 0) {
+                    printf("QuantLLM_API Error: Unsupported N dimension %lu!\n", N_PowerOf2);
+                    return cudaErrorUnknown;
+                }
+                Kernel_Ex<TilingConfig<4, 1, 8>, float>(stream,
+                                                        Weight1,
+                                                        Weight2,
+                                                        Scales,
+                                                        B,
+                                                        Reduction_Workspace,
+                                                        M_Global,
+                                                        N_Global,
+                                                        K_Global,
+                                                        Split_K);
+                break;
+        }
+        // Reduction for SplitK
+        dim3 GridDim((M_Global * N_Global) / REDUCTION_ELEMENT_PER_THREADBLOCK, 1, 1);
+        dim3 BlockDim(WARP_SIZE, 1, 1);
+        SplitK_Reduction<<<GridDim, BlockDim, 0, stream>>>(
+            C, Reduction_Workspace, M_Global, N_Global, Split_K);
+    }
+    return cudaGetLastError();
+}
+
+/*
+Computes FP6-FP16 GEMM (PyTorch interface).
+
+[Mathematical Formula]
+Standard definition of linear layer:    Out = In * trans(W), where In, Out, and W are stored in
+row-major. After Equivalent transformation    :    trans(Out) = W * trans(In). Note that we do not
+perform "transpose" during runtime, we instead interpret the In/Out as column-major matrices when
+calling our CUDA kernel.
+
+[Inputs]
+  _in_feats:  tensor of shape [B, IC];                  // half
+  _weights:   int tensor of shape [OC, IC // 16 * 3];   // 3 INT32 words contains 16 FP6 weights.
+  _scales:    tensor of shape [OC];                     // half
+  splitK:     splitting the MatMul problem along K dimension for higher GPU utilization, default 1.
+[Outputs]
+  _out_feats: tensor of shape [B, OC];                  // half
+*/
+torch::Tensor fp6_linear_forward_cuda(torch::Tensor _in_feats,
+                                      torch::Tensor _weights,
+                                      torch::Tensor _scales,
+                                      int splitK)
+{
+    int num_in_feats = _in_feats.size(0);
+    int num_in_channels = _in_feats.size(1);
+    int num_out_channels = _weights.size(0);
+    assert(num_in_channels % 64 == 0);
+    assert((num_in_channels / 16 * 3) ==
+           _weights.size(1));  // Making sure the K dimension is matched.
+    //
+    int M = num_out_channels;
+    int K = num_in_channels;
+    int N = num_in_feats;
+    // Input Tensors
+    auto weight1 = reinterpret_cast<const uint4*>(
+        _weights.data_ptr<int>());  // weights is [OC, IC] but in FP6.
+    auto weight2 = weight1 + num_in_channels * num_out_channels * 2 / 128;
+    auto in_feats = reinterpret_cast<const half*>(_in_feats.data_ptr<at::Half>());
+    auto scales = reinterpret_cast<const half*>(_scales.data_ptr<at::Half>());
+    // Output Tensors
+    auto options = torch::TensorOptions().dtype(_in_feats.dtype()).device(_in_feats.device());
+    at::Tensor _out_feats = torch::empty({num_in_feats, num_out_channels}, options);
+    auto out_feats = reinterpret_cast<half*>(_out_feats.data_ptr<at::Half>());
+
+    float* Reduction_Workspace = nullptr;
+    if (splitK != 1) {
+        auto options = torch::TensorOptions().dtype(torch::kFloat32).device(_in_feats.device());
+        at::Tensor _workspace = torch::empty({splitK, num_in_feats, num_out_channels}, options);
+        auto Reduction_Workspace = reinterpret_cast<float*>(
+            _out_feats.data_ptr<float>());  // Reduction_Workspace_Size = Split_K * M_Global *
+                                            // N_Global * sizeof(fp32)
+    }
+
+    fp6_linear_kernel(0,  // Using default stream here.
+                      weight1,
+                      weight2,
+                      scales,
+                      in_feats,
+                      out_feats,
+                      M,
+                      N,
+                      K,
+                      Reduction_Workspace,
+                      splitK);
+
+    return _out_feats;
+}
+
+/*
+ * Inputs:
+ * (1) unsigned char Weight_6bit [M*K*6/8]
+ * Outputs:
+ * (1) unsigned char Weight_2bit [M*K*2/8]
+ * (2) unsigned char Weight_4bit [M*K*4/8]
+ *
+ * Assumption: Weight_6bit, Weight_2bit, Weight_4bit all stored continuously in row-major.
+ * 8 FP6 = 6 Bytes
+ * 8 FP4 = 4 Bytes
+ * 8 FP2 = 2 Bytes
+ */
+
+/*
+ * Weight prepacking (Pytorch interface).
+ * [Input & Output]
+ *  fp6_tensor: int tensor of shape [OC, IC // 16 * 3];   // 3 INT32 words contains 16 FP6 weights.
+ * [Output]
+ *  packed_tensor: int tensor of shape [OC, IC // 16 * 3];
+ */
+torch::Tensor weight_matrix_prepacking_cpu(torch::Tensor fp6_tensor, size_t OC, size_t IC)
+{
+    assert((OC % 256 == 0) && (IC % 64 == 0));
+    assert((fp6_tensor.size(0) == OC) && (fp6_tensor.size(1) == IC / 16 * 3));
+    // auto packed_tensor = torch::empty_like(fp6_tensor);
+    // auto packed_tensor_ptr = reinterpret_cast<int*>(packed_tensor.data_ptr<int>());
+    auto fp6_tensor_ptr = reinterpret_cast<int*>(fp6_tensor.data_ptr<int>());
+    weight_matrix_prepacking(fp6_tensor_ptr, OC, IC);
+    return fp6_tensor;
+}

deepspeed/inference/v2/kernels/core_ops/cuda_linear/linear_kernels_cuda.h

@@ -0,0 +1,51 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+// This is a copy of FP6-LLM kernel code: https://arxiv.org/abs/2401.14112
+
+#ifndef DEEPSPEED_CUDA_LINEAR_FP6_LINEAR_CUH
+#define DEEPSPEED_CUDA_LINEAR_FP6_LINEAR_CUH
+
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <torch/extension.h>
+
+/*
+ * Computes FP6-FP16 GEMM (C++ interface).
+ */
+cudaError_t fp6_linear_kernel(cudaStream_t stream,
+                              const uint4* Weight1,
+                              const uint4* Weight2,
+                              const half* Scales,
+                              const half* B,
+                              half* C,
+                              const size_t M_Global,
+                              const size_t N_Global,
+                              const size_t K_Global,
+                              float* Reduction_Workspace,  // Reduction_Workspace_Size = Split_K *
+                                                           // M_Global * N_Global * sizeof(fp32)
+                              int Split_K);
+
+/*
+ * Computes FP6-FP16 GEMM (PyTorch interface).
+ */
+torch::Tensor fp6_linear_forward_cuda(torch::Tensor _in_feats,
+                                      torch::Tensor _weights,
+                                      torch::Tensor _scales,
+                                      int splitK = 1);
+
+/*
+ * In-place weight prepacking (C++ interface).
+ */
+void weight_matrix_prepacking(int* FP6Weights, size_t M, size_t K);
+
+/*
+ * Weight prepacking (Pytorch interface).
+ */
+torch::Tensor weight_matrix_prepacking_cpu(torch::Tensor fp6_tensor, size_t M, size_t K);
+
+#endif

deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .rms_norm import CUDARMSNorm
+from .rms_pre_norm import CUDARMSPreNorm

deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm.cpp

@@ -0,0 +1,123 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "rms_norm.h"
+
+#ifdef BF16_AVAILABLE
+#define DISPATCH_FOR_FLOAT(DTYPE, ...)                                  \
+    [&] {                                                               \
+        if (DTYPE == torch::kFloat32) {                                 \
+            using scalar_t = float;                                     \
+            return __VA_ARGS__();                                       \
+        } else if (DTYPE == torch::kFloat16) {                          \
+            using scalar_t = __half;                                    \
+            return __VA_ARGS__();                                       \
+        } else if (DTYPE == torch::kBFloat16) {                         \
+            using scalar_t = __nv_bfloat16;                             \
+            return __VA_ARGS__();                                       \
+        } else {                                                        \
+            TORCH_CHECK(false, "Unsupported dtype for BiasActivation"); \
+        }                                                               \
+    }()
+#else
+#define DISPATCH_FOR_FLOAT(DTYPE, ...)                                  \
+    [&] {                                                               \
+        if (DTYPE == torch::kFloat32) {                                 \
+            using scalar_t = float;                                     \
+            return __VA_ARGS__();                                       \
+        } else if (DTYPE == torch::kFloat16) {                          \
+            using scalar_t = __half;                                    \
+            return __VA_ARGS__();                                       \
+        } else {                                                        \
+            TORCH_CHECK(false, "Unsupported dtype for BiasActivation"); \
+        }                                                               \
+    }()
+#endif
+
+void rms_norm(torch::Tensor& norm_output,
+              torch::Tensor& norm_input,
+              torch::Tensor& gamma,
+              float epsilon)
+{
+    TORCH_CHECK(norm_output.scalar_type() == norm_input.scalar_type(),
+                "norm_output and norm_input should have the same data type");
+    TORCH_CHECK(norm_output.scalar_type() == gamma.scalar_type(),
+                "norm_output and gamma should have the same data type");
+
+    const int32_t rows = norm_input.size(0);
+    const int32_t cols = norm_input.size(1);
+
+    TORCH_CHECK(norm_output.size(0) == rows,
+                "norm_output and norm_input should have the same first dimension");
+    TORCH_CHECK(norm_output.size(1) == cols,
+                "norm_output and norm_input should have the same second dimension");
+
+    DISPATCH_FOR_FLOAT(norm_output.scalar_type(), [&] {
+        scalar_t* norm_output_ptr = reinterpret_cast<scalar_t*>(norm_output.data_ptr());
+        scalar_t* norm_input_ptr = reinterpret_cast<scalar_t*>(norm_input.data_ptr());
+        scalar_t* gamma_ptr = reinterpret_cast<scalar_t*>(gamma.data_ptr());
+        scalar_t* null_t = nullptr;
+
+        launch_rms_norm(norm_output_ptr,
+                        null_t,
+                        norm_input_ptr,
+                        null_t,
+                        gamma_ptr,
+                        epsilon,
+                        rows,
+                        cols,
+                        at::cuda::getCurrentCUDAStream());
+    });
+}
+
+void rms_pre_norm(torch::Tensor& norm_output,
+                  torch::Tensor& residual_output,
+                  torch::Tensor& norm_input,
+                  torch::Tensor& residual_input,
+                  torch::Tensor& gamma,
+                  float epsilon)
+{
+    TORCH_CHECK(norm_output.scalar_type() == norm_input.scalar_type(),
+                "norm_output and norm_input should have the same data type");
+    TORCH_CHECK(norm_output.scalar_type() == gamma.scalar_type(),
+                "norm_output and gamma should have the same data type");
+
+    const int32_t rows = norm_input.size(0);
+    const int32_t cols = norm_input.size(1);
+
+    TORCH_CHECK(norm_output.size(0) == rows,
+                "norm_output and norm_input should have the same first dimension");
+    TORCH_CHECK(norm_output.size(1) == cols,
+                "norm_output and norm_input should have the same second dimension");
+
+    TORCH_CHECK(residual_output.size(0) == rows,
+                "residual_output and norm_input should have the same first dimension");
+    TORCH_CHECK(residual_output.size(1) == cols,
+                "residual_output and norm_input should have the same second dimension");
+
+    TORCH_CHECK(residual_input.size(0) == rows,
+                "residual_input and norm_input should have the same first dimension");
+    TORCH_CHECK(residual_input.size(1) == cols,
+                "residual_input and norm_input should have the same second dimension");
+
+    DISPATCH_FOR_FLOAT(norm_output.scalar_type(), [&] {
+        scalar_t* norm_output_ptr = reinterpret_cast<scalar_t*>(norm_output.data_ptr());
+        scalar_t* residual_output_ptr = reinterpret_cast<scalar_t*>(residual_output.data_ptr());
+        const scalar_t* norm_input_ptr = reinterpret_cast<const scalar_t*>(norm_input.data_ptr());
+        const scalar_t* residual_input_ptr =
+            reinterpret_cast<const scalar_t*>(residual_input.data_ptr());
+        const scalar_t* gamma_ptr = reinterpret_cast<const scalar_t*>(gamma.data_ptr());
+
+        launch_rms_norm(norm_output_ptr,
+                        residual_output_ptr,
+                        norm_input_ptr,
+                        residual_input_ptr,
+                        gamma_ptr,
+                        epsilon,
+                        rows,
+                        cols,
+                        at::cuda::getCurrentCUDAStream());
+    });
+}

deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm.h

@@ -0,0 +1,33 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+#include "ds_kernel_utils.h"
+
+template <typename T>
+void launch_rms_norm(T* norm_output,
+                     T* res_output,
+                     const T* vals,
+                     const T* residual,
+                     const T* gamma,
+                     float epsilon,
+                     int rows,
+                     int elems_per_row,
+                     cudaStream_t stream);
+
+void rms_norm(torch::Tensor& norm_output,
+              torch::Tensor& norm_input,
+              torch::Tensor& gamma,
+              float epsilon);
+
+void rms_pre_norm(torch::Tensor& norm_output,
+                  torch::Tensor& residual_output,
+                  torch::Tensor& norm_input,
+                  torch::Tensor& residual_input,
+                  torch::Tensor& gamma,
+                  float epsilon);

deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm.py

@@ -0,0 +1,28 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from .rms_norm_base import CUDARMSNormBase
+
+
+class CUDARMSNorm(CUDARMSNormBase):
+    """
+    Floating point layer norm kernel for CUDA/RoCM.
+
+    Performs: z = ln(x)
+    """
+
+    def __call__(self, output_z: torch.Tensor, input_x: torch.Tensor, gamma: torch.Tensor) -> torch.Tensor:
+        """
+        output_z may alias input_x directly. All Tensors should have the same shape.
+
+        Parameters:
+            output_z (torch.Tensor): Output tensor.
+            input_x (torch.Tensor): Input tensor.
+            gamma (torch.Tensor): Gamma tensor.
+        """
+        self.inf_module.rms_norm(output_z, input_x, gamma, self.epsilon)
+        return output_z

deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm_base.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ... import DSKernelBase
+from ....inference_utils import elem_size
+from deepspeed.ops.op_builder import InferenceCoreBuilder
+
+
+class CUDARMSNormBase(DSKernelBase):
+    """
+    Base class for CUDA LN kernels. They all same the same validation logic,
+    so we can share it here.
+    """
+
+    supported_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+
+    def __init__(self, channels: int, fp_dtype: torch.dtype, epsilon: float = 1e-5):
+        """
+        Parameters:
+            channels (int): Number of channels in the input tensor. Must be divisible to align
+                to 16 bytes.
+            fp_dtype (torch.dtype): Data type for the input/output/gamma. Supported values
+                are torch.float16, torch.bfloat16, and torch.float32.
+        """
+        if fp_dtype not in CUDARMSNormBase.supported_dtypes:
+            raise ValueError("Unsupported data type: {}, supported_dtypes are {}".format(
+                fp_dtype, CUDARMSNormBase.supported_dtypes))
+
+        if elem_size(fp_dtype) * channels % 16 != 0:
+            raise ValueError("channels must be divisible by 16 bytes")
+
+        self.inf_module = InferenceCoreBuilder().load()
+        self.epsilon = epsilon

deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_norm_cuda.cu

@@ -0,0 +1,262 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "conversion_utils.h"
+#include "ds_kernel_utils.h"
+#include "memory_access_utils.h"
+#include "reduction_utils.h"
+
+namespace cg = cooperative_groups;
+using rop = reduce::ROpType;
+
+namespace rms {
+constexpr int granularity = 16;
+}  // namespace rms
+
+template <typename T, int UNROLL, int threadsPerGroup, int maxThreads>
+__global__ void rms_norm(T* output, const T* vals, const T* gamma, float epsilon, int elems_per_row)
+{
+    constexpr int T_per_load = rms::granularity / sizeof(T);
+
+    cg::thread_block tb = cg::this_thread_block();
+    cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);
+
+    // X-dimension of the block
+    const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) +
+                             (tb.thread_index().y * elems_per_row);
+    const int thread_offset = tb.thread_index().x * T_per_load;
+    const int base_offset = block_offset + thread_offset;
+    const int stride = blockDim.x * T_per_load;
+
+    float var_sum = reduce::init<rop::Add, float>();
+
+    const T* input_base = vals + base_offset;
+
+    T local_buffer[UNROLL * T_per_load];
+
+#pragma unroll
+    for (int i = 0; i < UNROLL; i++) {
+        T* iteration_buffer = local_buffer + (i * T_per_load);
+
+        mem_access::load_global<rms::granularity>(iteration_buffer,
+                                                  input_base + (i * stride),
+                                                  thread_offset + (i * stride) < elems_per_row);
+
+#pragma unroll
+        for (int j = 0; j < T_per_load; j++) {
+            float up_cast = conversion::to<float>(iteration_buffer[j]);
+            float sq_val = up_cast * up_cast;
+            var_sum = reduce::element<rop::Add, float>(var_sum, sq_val);
+        }
+    }
+
+    reduce::partitioned_block<rop::Add, threadsPerGroup>(tb, warp, var_sum);
+    const float var = var_sum / elems_per_row;
+    const T denom = conversion::to<T>(__frsqrt_rn(var + epsilon));
+
+    T* block_output = output + block_offset;
+
+#pragma unroll
+    for (int i = 0; i < UNROLL; i++) {
+        T* iteration_buffer = local_buffer + (i * T_per_load);
+        const int iter_idx = i * stride + thread_offset;
+        const bool do_loads = (iter_idx < elems_per_row);
+
+        T gamma_local[T_per_load];
+
+        mem_access::load_global<rms::granularity>(gamma_local, gamma + iter_idx, do_loads);
+
+#pragma unroll
+        for (int j = 0; j < T_per_load; j++) {
+            iteration_buffer[j] *= denom;
+            iteration_buffer[j] *= gamma_local[j];
+        }
+
+        if (do_loads) {
+            mem_access::store_global<rms::granularity>(block_output + iter_idx, iteration_buffer);
+        }
+    }
+}
+
+template <typename T, int UNROLL, int threadsPerGroup, int maxThreads>
+__global__ void pre_rms_norm(T* output,
+                             T* res_out,
+                             const T* vals,
+                             const T* residual,
+                             const T* gamma,
+                             float epsilon,
+                             int elems_per_row)
+{
+    constexpr int T_per_load = rms::granularity / sizeof(T);
+
+    cg::thread_block tb = cg::this_thread_block();
+    cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);
+
+    // X-dimension of the block
+    const int block_offset = (tb.group_index().x * (maxThreads / threadsPerGroup) * elems_per_row) +
+                             (tb.thread_index().y * elems_per_row);
+    const int thread_offset = tb.thread_index().x * T_per_load;
+    const int base_offset = block_offset + thread_offset;
+    const int stride = blockDim.x * T_per_load;
+
+    float var_sum = reduce::init<rop::Add, float>();
+
+    const T* input_base = vals + base_offset;
+    const T* residual_base = residual + base_offset;
+    T* res_output = res_out + base_offset;
+
+    T local_buffer[UNROLL * T_per_load];
+
+#pragma unroll
+    for (int i = 0; i < UNROLL; i++) {
+        T* iteration_buffer = local_buffer + (i * T_per_load);
+        T residual_buffer[T_per_load];
+
+        const int iter_offset = i * stride + thread_offset;
+        const bool do_loads = (iter_offset < elems_per_row);
+
+        mem_access::load_global<rms::granularity>(
+            iteration_buffer, input_base + (i * stride), do_loads);
+        mem_access::load_global<rms::granularity>(
+            residual_buffer, residual_base + (i * stride), do_loads);
+
+#pragma unroll
+        for (int j = 0; j < T_per_load; j++) {
+            iteration_buffer[j] += residual_buffer[j];
+            float vals_up_cast = conversion::to<float>(iteration_buffer[j]);
+
+            var_sum = reduce::element<rop::Add, float>(var_sum, vals_up_cast * vals_up_cast);
+        }
+
+        if (do_loads) {
+            mem_access::store_global<rms::granularity>(res_output + i * stride, iteration_buffer);
+        }
+    }
+
+    reduce::partitioned_block<rop::Add, threadsPerGroup>(tb, warp, var_sum);
+    const float var = var_sum / elems_per_row;
+    const T denom = conversion::to<T>(__frsqrt_rn(var + epsilon));
+
+    T* block_output = output + block_offset;
+
+#pragma unroll
+    for (int i = 0; i < UNROLL; i++) {
+        T* iteration_buffer = local_buffer + (i * T_per_load);
+        const int iter_idx = i * stride + thread_offset;
+        const bool do_loads = (iter_idx < elems_per_row);
+
+        T gamma_local[T_per_load];
+
+        mem_access::load_global<rms::granularity>(gamma_local, gamma + iter_idx, do_loads);
+
+#pragma unroll
+        for (int j = 0; j < T_per_load; j++) {
+            iteration_buffer[j] *= denom;
+            iteration_buffer[j] *= gamma_local[j];
+        }
+
+        if (do_loads) {
+            mem_access::store_global<rms::granularity>(block_output + iter_idx, iteration_buffer);
+        }
+    }
+}
+
+#define LAUNCH_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \
+    rms_norm<T, UNROLL, threadsPerGroup, maxThreads>         \
+        <<<grid, block, 0, stream>>>(norm_output, vals, gamma, epsilon, elems_per_row);
+
+#define LAUNCH_PRE_RMS_NORM(UNROLL, threadsPerGroup, maxThreads)                      \
+    pre_rms_norm<T, UNROLL, threadsPerGroup, maxThreads><<<grid, block, 0, stream>>>( \
+        norm_output, res_output, vals, residual, gamma, epsilon, elems_per_row);
+
+#define LAUNCH_ALL_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \
+    if (pre_norm) {                                              \
+        LAUNCH_PRE_RMS_NORM(UNROLL, threadsPerGroup, maxThreads) \
+    } else {                                                     \
+        LAUNCH_RMS_NORM(UNROLL, threadsPerGroup, maxThreads)     \
+    }
+
+template <typename T>
+void launch_rms_norm(T* norm_output,
+                     T* res_output,
+                     const T* vals,
+                     const T* residual,
+                     const T* gamma,
+                     float epsilon,
+                     int rows,
+                     int elems_per_row,
+                     cudaStream_t stream)
+{
+    // 8 for __half, 4 for float
+    constexpr int T_per_load = rms::granularity / sizeof(T);
+    constexpr int maxThreads = 256;
+    constexpr int internalUnroll = sizeof(T) == 4 ? 4 : 2;
+
+    const bool is_subblock_schedule = (elems_per_row <= 128) ? true : false;
+    const int h_per_step = is_subblock_schedule ? T_per_load : T_per_load * internalUnroll;
+
+    // Scheduling concern: may be slightly faster for some inputs to assign multiple stages of
+    // warp-sized blocks rather than stepping up to 64/96 threads
+    const int one_step_threads = next_pow2((elems_per_row + h_per_step - 1) / h_per_step);
+    const int threads_per_group = (one_step_threads < maxThreads) ? one_step_threads : maxThreads;
+
+    const int groups_per_block_max =
+        is_subblock_schedule ? (maxThreads + threads_per_group - 1) / threads_per_group : 1;
+    const int groups_per_block = (rows < groups_per_block_max) ? rows : groups_per_block_max;
+    const int groups_launch = (groups_per_block + rows - 1) / groups_per_block;
+
+    dim3 block(threads_per_group, groups_per_block);
+    dim3 grid(groups_launch);
+
+    const int elems_per_step = threads_per_group * h_per_step;
+    const int external_unRoll = (elems_per_row + elems_per_step - 1) / elems_per_step;
+
+    bool pre_norm = (residual == nullptr) ? false : true;
+
+    if (is_subblock_schedule) {
+        // <=128
+        if (threads_per_group == 1) {
+            LAUNCH_ALL_RMS_NORM(1, 1, maxThreads);
+        } else if (threads_per_group == 2) {
+            LAUNCH_ALL_RMS_NORM(1, 2, maxThreads);
+        } else if (threads_per_group == 4) {
+            LAUNCH_ALL_RMS_NORM(1, 4, maxThreads);
+        } else if (threads_per_group == 8) {
+            LAUNCH_ALL_RMS_NORM(1, 8, maxThreads);
+        } else if (threads_per_group == 16) {
+            LAUNCH_ALL_RMS_NORM(1, 16, maxThreads);
+        }
+    } else if (external_unRoll == 1) {
+        // 129 - 4096 elems
+        // (this can launch with 1-7 warps as well)
+        LAUNCH_ALL_RMS_NORM(1 * internalUnroll, maxThreads, maxThreads);
+    } else if (external_unRoll == 2) {
+        // 4097 - 8192 elems
+        LAUNCH_ALL_RMS_NORM(2 * internalUnroll, maxThreads, maxThreads);
+    } else if (external_unRoll == 3) {
+        // 8193 - 12288 elems
+        LAUNCH_ALL_RMS_NORM(3 * internalUnroll, maxThreads, maxThreads);
+    } else if (external_unRoll == 4) {
+        // 12289 - 16384 elems
+        LAUNCH_ALL_RMS_NORM(4 * internalUnroll, maxThreads, maxThreads);
+    }
+}
+
+#define INSTANTIATE_LAUNCH_RMS_NORM(T)                  \
+    template void launch_rms_norm<T>(T * norm_output,   \
+                                     T * res_output,    \
+                                     const T* vals,     \
+                                     const T* residual, \
+                                     const T* gamma,    \
+                                     float epsilon,     \
+                                     int rows,          \
+                                     int elems_per_row, \
+                                     cudaStream_t stream);
+
+INSTANTIATE_LAUNCH_RMS_NORM(float)
+INSTANTIATE_LAUNCH_RMS_NORM(__half)
+#ifdef BF16_AVAILABLE
+INSTANTIATE_LAUNCH_RMS_NORM(__nv_bfloat16)
+#endif

deepspeed/inference/v2/kernels/core_ops/cuda_rms_norm/rms_pre_norm.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Tuple
+
+import torch
+
+from .rms_norm_base import CUDARMSNormBase
+
+
+class CUDARMSPreNorm(CUDARMSNormBase):
+    """
+    Floating point pre-LayerNorm kernel for CUDA/RoCM.
+
+    Performs: z_res = x_res + y_hid
+              z_hid = ln(z_hid)
+    """
+
+    def __call__(self, z_res: torch.Tensor, z_hid: torch.Tensor, x_res: torch.Tensor, y_hid: torch.Tensor,
+                 gamma: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        z_res can alias x_res. All non-parameter input/output tensors
+        must have the same shape. z_hid can alias y_hid.
+
+        Parameters:
+            z_res (torch.Tensor): Output residual.
+            z_hid (torch.Tensor): Output hidden states.
+            x_res (torch.Tensor): Input residual.
+            y_hid (torch.Tensor): Input hidden states.
+            gamma (torch.Tensor): Gamma tensor.
+            beta (torch.Tensor): Beta tensor.
+
+        Returns:
+            output (torch.Tensor): Output tensor.
+        """
+        self.inf_module.rms_pre_norm(z_hid, z_res, y_hid, x_res, gamma, self.epsilon)
+        return z_res, z_hid

deepspeed/inference/v2/kernels/core_ops/gated_activations/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .gated_activation import *

deepspeed/inference/v2/kernels/core_ops/gated_activations/gated_activation.py

@@ -0,0 +1,65 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional
+
+import torch
+
+from ... import DSKernelBase
+from ....inference_utils import ActivationType, elem_size
+from deepspeed.ops.op_builder import InferenceCoreBuilder
+
+
+class CUDAGatedActivation(DSKernelBase):
+    """
+    CUDA implementation of gated activation kernel. This kernel assumes that the input
+    tensor has gate and activation values in adjacent channels. The output tensor should
+    have half the dimensionality of the input tensor.
+    """
+
+    supported_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+    supported_act_fns = [ActivationType.GEGLU, ActivationType.ReGLU, ActivationType.SiGLU]
+
+    def __init__(self, channels: int, fp_dtype: torch.dtype, act_fn: ActivationType) -> None:
+        """
+        Compile and validate for the gated activation function.
+
+        Args:
+            channels (int): Number of columns in the output tensor. Must be divisible to align
+                to 8 bytes.
+            fp_dtype (torch.dtype): Data type for the input/output/gamma. Supported values
+                are torch.float16, torch.bfloat16, and torch.float32.
+            act_fn (ActivationType): Activation function to use. Only GEGLU is supported.
+        """
+        if fp_dtype not in CUDAGatedActivation.supported_dtypes:
+            raise ValueError("Unsupported data type: {}, supported_dtypes are {}".format(
+                fp_dtype, CUDAGatedActivation.supported_dtypes))
+
+        act_fn = ActivationType(act_fn)
+        if act_fn not in CUDAGatedActivation.supported_act_fns:
+            raise ValueError("Unsupported activation function: {}, supported_act_fns are {}".format(
+                act_fn, CUDAGatedActivation.supported_act_fns))
+
+        if elem_size(fp_dtype) * channels % 8 != 0:
+            raise ValueError("Channels must be divisible by 16 bytes")
+
+        if elem_size(fp_dtype) * channels > 98304:
+            raise ValueError(
+                "Kernel only compiled to support 98304 bytes per row, please file an issue if your model requires more."
+            )
+
+        self.inf_module = InferenceCoreBuilder().load()
+        self.act_fn = act_fn
+        self.kernel = self.inf_module.gated_activation
+
+    def __call__(self, output: torch.Tensor, input: torch.Tensor, bias: Optional[torch.Tensor] = None) -> None:
+        """
+        Performs gated activation on the input tensor, writing the result to the output tensor.
+
+        Args:
+            output (torch.Tensor): Output tensor. Can be of [T, C // 2] or [B, S, C // 2]
+            input (torch.Tensor): Input tensor. Can be of [T, C] or [B, S, C]
+        """
+        self.kernel(output, input, bias, self.act_fn.value)

deepspeed/inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels.cpp

@@ -0,0 +1,72 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "gated_activation_kernels.h"
+
+#ifdef BF16_AVAILABLE
+#define DISPATCH_FOR_FLOAT(DTYPE, ...)                                  \
+    [&] {                                                               \
+        if (DTYPE == torch::kFloat32) {                                 \
+            using scalar_t = float;                                     \
+            return __VA_ARGS__();                                       \
+        } else if (DTYPE == torch::kFloat16) {                          \
+            using scalar_t = __half;                                    \
+            return __VA_ARGS__();                                       \
+        } else if (DTYPE == torch::kBFloat16) {                         \
+            using scalar_t = __nv_bfloat16;                             \
+            return __VA_ARGS__();                                       \
+        } else {                                                        \
+            TORCH_CHECK(false, "Unsupported dtype for BiasActivation"); \
+        }                                                               \
+    }()
+#else
+#define DISPATCH_FOR_FLOAT(DTYPE, ...)                                  \
+    [&] {                                                               \
+        if (DTYPE == torch::kFloat32) {                                 \
+            using scalar_t = float;                                     \
+            return __VA_ARGS__();                                       \
+        } else if (DTYPE == torch::kFloat16) {                          \
+            using scalar_t = __half;                                    \
+            return __VA_ARGS__();                                       \
+        } else {                                                        \
+            TORCH_CHECK(false, "Unsupported dtype for BiasActivation"); \
+        }                                                               \
+    }()
+#endif
+
+void ds_gated_activation(at::Tensor& output,
+                         at::Tensor& input,
+                         c10::optional<torch::Tensor>& bias,
+                         int activation_type_raw)
+{
+    bool ragged_input = input.dim() == 2;
+
+    const ActivationType activation_type = static_cast<ActivationType>(activation_type_raw);
+
+    const int rows = ragged_input ? input.size(0) : input.size(0) * input.size(1);
+    const int cols = ragged_input ? input.size(1) : input.size(2);
+
+    DISPATCH_FOR_FLOAT(input.scalar_type(), [&] {
+        scalar_t* bias_ptr = nullptr;
+        if (bias.has_value()) {
+            TORCH_CHECK(bias.value().scalar_type() == input.scalar_type(),
+                        "Bias type must match input type");
+            TORCH_CHECK(bias.value().numel() == cols,
+                        "Bias must have the same number of elements as the input channels");
+            bias_ptr = reinterpret_cast<scalar_t*>(bias.value().data_ptr());
+        }
+
+        scalar_t* output_ptr = reinterpret_cast<scalar_t*>(output.data_ptr());
+        const scalar_t* input_ptr = reinterpret_cast<const scalar_t*>(input.data_ptr());
+
+        launch_gated_activation(output_ptr,
+                                input_ptr,
+                                bias_ptr,
+                                rows,
+                                cols,
+                                activation_type,
+                                c10::cuda::getCurrentCUDAStream());
+    });
+}

deepspeed/inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels.h

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+#include "activation_type.h"
+#include "ds_kernel_utils.h"
+
+template <typename T>
+void launch_gated_activation(T* output,
+                             const T* vals,
+                             const T* bias,
+                             int rows,
+                             int cols,
+                             ActivationType activation_type,
+                             cudaStream_t stream);
+
+void ds_gated_activation(at::Tensor& output,
+                         at::Tensor& input,
+                         c10::optional<torch::Tensor>& bias,
+                         int activation_type_raw);

deepspeed/inference/v2/kernels/core_ops/gated_activations/gated_activation_kernels_cuda.cu

@@ -0,0 +1,172 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include <stdexcept>
+#include "activation_type.h"
+#include "conversion_utils.h"
+#include "ds_kernel_utils.h"
+#include "memory_access_utils.h"
+
+namespace cg = cooperative_groups;
+
+namespace gated_act {
+
+constexpr int access_size = 16;
+constexpr int threads = 1024;
+
+template <ActivationType ActType>
+DS_D_INLINE float gated_act_fn(float x, float y);
+
+template <>
+DS_D_INLINE float gated_act_fn<ActivationType::GEGLU>(float x, float y)
+{
+    constexpr float sqrt_param = 0.79788456080286535587989211986876f;
+    constexpr float mul_param = 0.044715;
+    return y * x * 0.5f * (1.0f + tanhf(sqrt_param * (x + mul_param * x * x * x)));
+}
+
+template <>
+DS_D_INLINE float gated_act_fn<ActivationType::ReGLU>(float x, float y)
+{
+    return y * (x > 0.0f ? x : 0.0f);
+}
+
+template <>
+DS_D_INLINE float gated_act_fn<ActivationType::SiGLU>(float x, float y)
+{
+    return y * (x / (1.0f + expf(-x)));
+}
+
+}  // namespace gated_act
+
+template <typename T, ActivationType ActType, int loopUnroll>
+__global__ void gated_activation_kernel(T* output,
+                                        const T* input,
+                                        const T* bias,
+                                        int rows,
+                                        int cols)
+{
+    constexpr int read_vector = gated_act::access_size / sizeof(T);
+    constexpr int write_vector = read_vector / 2;
+
+    const int row = blockIdx.x;
+    const int col = threadIdx.x * read_vector;
+
+    const T* input_row = input + row * cols;
+    T* output_row = output + row * cols / 2;
+
+#pragma unroll
+    for (int i = 0; i < loopUnroll; i++) {
+        T read[read_vector];
+        T bias_read[read_vector];
+        T store[write_vector];
+
+        const int read_offset = col + gated_act::threads * read_vector * i;
+        const int write_offset = col / 2 + gated_act::threads * write_vector * i;
+
+        if (i != loopUnroll - 1 || read_offset < cols) {
+            mem_access::load_global<gated_act::access_size>(read, input_row + read_offset);
+            mem_access::load_global<gated_act::access_size>(
+                bias_read, bias + read_offset, bias != nullptr);
+
+            for (int j = 0; j < write_vector; j++) {
+                float g_val =
+                    conversion::to<float>(read[j * 2]) + conversion::to<float>(bias_read[j * 2]);
+                float a_val = conversion::to<float>(read[j * 2 + 1]) +
+                              conversion::to<float>(bias_read[j * 2 + 1]);
+
+                float act_val = gated_act::gated_act_fn<ActType>(g_val, a_val);
+                store[j] = conversion::to<T>(act_val);
+            }
+
+            mem_access::store_global<gated_act::access_size / 2>(output_row + write_offset, store);
+        }
+    }
+}
+
+#define DISPATCH_UNROLL(unroll_val)                 \
+    gated_activation_kernel<T, ActType, unroll_val> \
+        <<<grid, block, 0, stream>>>(output, input, bias, rows, cols);
+
+template <typename T, ActivationType ActType>
+void launch_gated_activation_impl(T* output,
+                                  const T* input,
+                                  const T* bias,
+                                  int rows,
+                                  int cols,
+                                  cudaStream_t stream)
+{
+    constexpr int read_vector = gated_act::access_size / sizeof(T);
+    constexpr int cols_per_unroll = gated_act::threads * read_vector;
+    const int req_threads = (cols + read_vector - 1) / read_vector;
+    const int threads = std::min(req_threads, gated_act::threads);
+
+    const dim3 grid(rows);
+    const dim3 block(threads);
+    const int unroll = (cols + cols_per_unroll - 1) / cols_per_unroll;
+
+    if (unroll == 1) {
+        DISPATCH_UNROLL(1);
+    } else if (unroll == 2) {
+        DISPATCH_UNROLL(2);
+    } else if (unroll == 3) {
+        DISPATCH_UNROLL(3);
+    } else if (unroll == 4) {
+        DISPATCH_UNROLL(4);
+    } else if (unroll == 5) {
+        DISPATCH_UNROLL(5);
+    } else if (unroll == 6) {
+        DISPATCH_UNROLL(6);
+    } else if (unroll == 7) {
+        DISPATCH_UNROLL(7);
+    } else {
+        // TODO: provide a kernel with an outer loop to handle larger columns.
+        throw std::runtime_error(
+            "Called with more columns than supported, please report this bug and this limit will "
+            "be increased.");
+    }
+}
+
+template <typename T>
+void launch_gated_activation(T* output,
+                             const T* input,
+                             const T* bias,
+                             int rows,
+                             int cols,
+                             ActivationType act_type,
+                             cudaStream_t stream)
+{
+    switch (act_type) {
+        case ActivationType::GEGLU:
+            launch_gated_activation_impl<T, ActivationType::GEGLU>(
+                output, input, bias, rows, cols, stream);
+            break;
+        case ActivationType::ReGLU:
+            launch_gated_activation_impl<T, ActivationType::ReGLU>(
+                output, input, bias, rows, cols, stream);
+            break;
+        case ActivationType::SiGLU:
+            launch_gated_activation_impl<T, ActivationType::SiGLU>(
+                output, input, bias, rows, cols, stream);
+            break;
+        default: throw std::runtime_error("Unsupported activation type");
+    }
+}
+
+#define INSTANTIATE_FOR_TYPE(T)                                       \
+    template void launch_gated_activation<T>(T * output,              \
+                                             const T* input,          \
+                                             const T* bias,           \
+                                             int rows,                \
+                                             int cols,                \
+                                             ActivationType act_type, \
+                                             cudaStream_t stream);
+
+INSTANTIATE_FOR_TYPE(float)
+INSTANTIATE_FOR_TYPE(__half)
+
+#ifdef BF16_AVAILABLE
+INSTANTIATE_FOR_TYPE(__nv_bfloat16)
+#endif

deepspeed/inference/v2/kernels/cutlass_ops/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

deepspeed/inference/v2/kernels/cutlass_ops/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .mixed_gemm import *
+from .moe_gemm import *

deepspeed/inference/v2/kernels/cutlass_ops/cutlass_ops.cpp

@@ -0,0 +1,19 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include <torch/extension.h>
+
+#include "mixed_gemm.h"
+#include "moe_gemm.h"
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    // mixed_gemm.h
+    m.def("mixed_gemm", &mixed_gemm, "Mixed-precision GEMM");
+
+    // moe_gemm.h
+    m.def("moe_gemm", &moe_gemm, "MultiGEMM for MoE (16-bit weights)");
+    m.def("mixed_moe_gemm", &mixed_moe_gemm, "MultiGEMM for MoE (4-bit/8-bit weights)");
+}

deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .mixed_gemm import *

deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm.cu

@@ -0,0 +1,93 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include <c10/cuda/CUDAStream.h>
+#include "mixed_gemm.h"
+#include "mixed_gemm_api.h"
+#include "weight_variant.h"
+
+// Switch helpers inspired by
+// https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
+// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h
+
+#define ACT_DTYPE_SWITCH(COND, ...)                \
+    [&] {                                          \
+        if (COND) {                                \
+            using ActivationDtype = __half;        \
+            return __VA_ARGS__();                  \
+        } else {                                   \
+            using ActivationDtype = __nv_bfloat16; \
+            return __VA_ARGS__();                  \
+        }                                          \
+    }()
+
+#define WEIGHT_VARIANT_SWITCH(COND, ...)                            \
+    [&] {                                                           \
+        if (COND) {                                                 \
+            constexpr WeightVariant WVariant = WeightVariant::kFP8; \
+            return __VA_ARGS__();                                   \
+        } else {                                                    \
+            constexpr WeightVariant WVariant = WeightVariant::kFP4; \
+            return __VA_ARGS__();                                   \
+        }                                                           \
+    }()
+
+void mixed_gemm(at::Tensor& output,
+                at::Tensor& hidden_states,
+                at::Tensor& weight,
+                at::Tensor& scales,
+                c10::optional<at::Tensor>& bias,
+                int num_bits,
+                int activation_raw)
+{
+    TORCH_CHECK(output.dtype() == hidden_states.dtype(),
+                "Output and hidden states must have the same dtype");
+    TORCH_CHECK(num_bits == 4 || num_bits == 8, "Data width must be 4 or 8");
+    TORCH_CHECK(output.size(0) == hidden_states.size(0), "Token dimension mismatch");
+
+    int32_t m = output.size(0);
+    int32_t k = hidden_states.size(1);
+    int32_t n = weight.size(1);
+
+    TORCH_CHECK(weight.size(0) == k, "Weight dimension mismatch");
+
+    ACT_DTYPE_SWITCH(hidden_states.dtype() == torch::kFloat16, [&] {
+        WEIGHT_VARIANT_SWITCH(num_bits == 8, [&] {
+            fastertransformer::CutlassFpAIntBGemmRunner<ActivationDtype, WVariant> runner =
+                *MixedGemmContext<ActivationDtype, WVariant>::Instance().GeMM_Runner();
+
+            ActivationType activation_type = (ActivationType)activation_raw;
+            if (!bias.has_value() && activation_type == ActivationType::IDENTITY) {
+                runner.gemm((ActivationDtype*)hidden_states.data_ptr(),
+                            (const char*)weight.data_ptr(),
+                            (ActivationDtype*)scales.data_ptr(),
+                            (ActivationDtype*)output.data_ptr(),
+                            m,
+                            n,
+                            k,
+                            nullptr,
+                            0,
+                            at::cuda::getCurrentCUDAStream());
+                return;
+            } else {
+                ActivationDtype* bias_ptr = nullptr;
+                if (bias.has_value()) { bias_ptr = (ActivationDtype*)bias.value().data_ptr(); }
+                runner.gemm_bias_act((ActivationDtype*)hidden_states.data_ptr(),
+                                     (char*)weight.data_ptr(),
+                                     (ActivationDtype*)scales.data_ptr(),
+                                     bias_ptr,
+                                     (ActivationDtype*)output.data_ptr(),
+                                     m,
+                                     n,
+                                     k,
+                                     activation_type,
+                                     nullptr,
+                                     0,
+                                     at::cuda::getCurrentCUDAStream());
+                return;
+            }
+        });
+    });
+}

deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm.h

@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <torch/extension.h>
+
+void mixed_gemm(at::Tensor& output,
+                at::Tensor& hidden_states,
+                at::Tensor& weight,
+                at::Tensor& scales,
+                c10::optional<at::Tensor>& bias,
+                int num_bits,
+                int activation_raw);

deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm.py

@@ -0,0 +1,64 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ... import DSKernelBase
+from ....inference_utils import ActivationType, DtypeEnum
+from deepspeed.ops.op_builder import InferenceCutlassBuilder
+
+from typing import Optional
+
+
+class MixedGEMM(DSKernelBase):
+    """
+    CUTLASS implementation of MoE GEMM.
+    """
+
+    supported_dtypes = [DtypeEnum.fp16, DtypeEnum.bf16]
+    supported_act_fns = [ActivationType.GELU, ActivationType.SILU, ActivationType.RELU, ActivationType.IDENTITY]
+
+    def __init__(self, fp_dtype: DtypeEnum, act_fn: ActivationType, num_bits: int) -> None:
+
+        if not isinstance(fp_dtype, DtypeEnum):
+            fp_dtype = DtypeEnum(fp_dtype)
+
+        if fp_dtype not in MixedGEMM.supported_dtypes:
+            raise ValueError("Unsupported data type: {}, supported_dtypes are {}".format(
+                fp_dtype, MixedGEMM.supported_dtypes))
+
+        if act_fn not in MixedGEMM.supported_act_fns:
+            raise ValueError("Unsupported activation function: {}, supported_act_fns are {}".format(
+                act_fn, MixedGEMM.supported_act_fns))
+
+        if num_bits != 4 and num_bits != 8:
+            raise ValueError("Unsupported num_bits: {}, supported num_bits are 4 and 8".format(num_bits))
+
+        inf_module = InferenceCutlassBuilder().load()
+        self.num_bits = num_bits
+        self.kernel = inf_module.moe_gemm
+        self.act_fn = act_fn
+
+    def __call__(self,
+                 output: torch.Tensor,
+                 hidden_states: torch.Tensor,
+                 weights: torch.Tensor,
+                 scales: torch.Tensor,
+                 biases: Optional[torch.Tensor] = None) -> None:
+        """
+            Performs a MoE GEMM. Note that the stride between token inputs must be even (the distance between byte 1 of token 0 and token 1 must be the same as the distance between byte 1 of token 1 and token 2).
+
+            Arguments:
+                output (torch.Tensor): The output of the MoE GEMM of shape [n_tokens, out_neurons].
+                hidden_states (torch.Tensor): The direct input for the MoE GEMM of shape [n_tokens, in_neurons].
+                weights (torch.Tensor): The weights of shape [in_neurons, out_neurons]. These weights must be contiguous.
+                scales (torch.Tensor): The scales of shape [out_neurons]. These scales must be contiguous.
+                biases (torch.Tensor): The biases of shape [out_neurons]. These biases must be contiguous.
+
+            Returns:
+                output
+            """
+        self.kernel(output, hidden_states, weights, biases, self.num_bits, self.act_fn)
+        return output

deepspeed/inference/v2/kernels/cutlass_ops/mixed_gemm/mixed_gemm_api.h

@@ -0,0 +1,57 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "activation_type.h"
+#include "weight_variant.h"
+
+namespace fastertransformer {
+
+template <typename T, WeightVariant V>
+class CutlassFpAIntBGemmRunner {
+public:
+    void gemm(const T* A,
+              const char* B,
+              const T* weight_scales,
+              T* C,
+              int m,
+              int n,
+              int k,
+              char* workspace_ptr,
+              const size_t workspace_bytes,
+              cudaStream_t stream);
+
+    void gemm_bias_act(const T* A,
+                       const char* B,
+                       const T* weight_scales,
+                       const T* biases,
+                       T* C,
+                       int m,
+                       int n,
+                       int k,
+                       ActivationType activation_type,
+                       char* workspace_ptr,
+                       const size_t workspace_bytes,
+                       cudaStream_t stream);
+};
+
+}  // namespace fastertransformer
+
+template <typename T, WeightVariant V>
+class MixedGemmContext {
+public:
+    MixedGemmContext() { _runner = new fastertransformer::CutlassFpAIntBGemmRunner<T, V>(); }
+
+    virtual ~MixedGemmContext() { delete _runner; }
+
+    static MixedGemmContext& Instance()
+    {
+        static MixedGemmContext _ctx;
+        return _ctx;
+    }
+
+    fastertransformer::CutlassFpAIntBGemmRunner<T, V>* GeMM_Runner() const { return _runner; }
+
+    fastertransformer::CutlassFpAIntBGemmRunner<T, V>* _runner;
+};

deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .mixed_moe_gemm import *
+from .moe_gemm import *

deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/mixed_moe_gemm.py

@@ -0,0 +1,67 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ... import DSKernelBase
+from ....inference_utils import ActivationType, DtypeEnum
+from deepspeed.ops.op_builder import InferenceCutlassBuilder
+
+from typing import Optional
+
+
+class MixedMoEGEMM(DSKernelBase):
+    """
+    CUTLASS implementation of MoE GEMM.
+    """
+
+    supported_dtypes = [DtypeEnum.fp16, DtypeEnum.bf16]
+    supported_act_fns = [ActivationType.GELU, ActivationType.SILU, ActivationType.RELU, ActivationType.IDENTITY]
+
+    def __init__(self, fp_dtype: DtypeEnum, act_fn: ActivationType, num_bits: int) -> None:
+
+        if not isinstance(fp_dtype, DtypeEnum):
+            fp_dtype = DtypeEnum(fp_dtype)
+
+        if fp_dtype not in MixedMoEGEMM.supported_dtypes:
+            raise ValueError("Unsupported data type: {}, supported_dtypes are {}".format(
+                fp_dtype, MixedMoEGEMM.supported_dtypes))
+
+        if act_fn not in MixedMoEGEMM.supported_act_fns:
+            raise ValueError("Unsupported activation function: {}, supported_act_fns are {}".format(
+                act_fn, MixedMoEGEMM.supported_act_fns))
+
+        if num_bits != 4 and num_bits != 8:
+            raise ValueError("Unsupported num_bits: {}, supported num_bits are 4 and 8".format(num_bits))
+
+        inf_module = InferenceCutlassBuilder().load()
+        self.num_bits = num_bits
+        self.kernel = inf_module.moe_gemm
+        self.act_fn = act_fn
+
+    def __call__(self,
+                 ordered_output: torch.Tensor,
+                 ordered_input: torch.Tensor,
+                 weights: torch.Tensor,
+                 scales: torch.Tensor,
+                 total_rows_before_expert: torch.Tensor,
+                 biases: Optional[torch.Tensor] = None) -> None:
+        """
+            Performs a MoE GEMM. Note that the stride between token inputs must be even (the distance between byte 1 of token 0 and token 1 must be the same as the distance between byte 1 of token 1 and token 2).
+
+            Arguments:
+                ordered_output (torch.Tensor): The output of the MoE GEMM of shape [n_tokens, out_neurons].
+                ordered_input (torch.Tensor): The direct input for the MoE GEMM of shape [n_tokens, in_neurons].
+                weights (torch.Tensor): The weights of shape [n_experts, in_neurons, out_neurons]. These weights must be contiguous.
+                scales (torch.Tensor): The scales of shape [n_experts, out_neurons]. These scales must be contiguous.
+                total_rows_before_expert (torch.Tensor): The total number of rows before each expert of shape [n_experts].
+                biases (torch.Tensor): The biases of shape [n_experts, out_neurons]. These biases must be contiguous.
+
+            Returns:
+                ordered_output
+            """
+        self.kernel(ordered_output, ordered_input, weights, scales, biases, total_rows_before_expert, self.num_bits,
+                    self.act_fn)
+        return ordered_output

deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm.cu

@@ -0,0 +1,175 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include <c10/cuda/CUDAStream.h>
+#include "moe_gemm.h"
+#include "moe_gemm_api.h"
+#include "weight_variant.h"
+
+// Switch helpers inspired by
+// https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
+// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h
+
+#define HIDDEN_DTYPE_SWITCH(COND, ...)                               \
+    [&] {                                                            \
+        if (COND) {                                                  \
+            using ActivationDtype = __half;                          \
+            constexpr WeightVariant WVariant = WeightVariant::kFP16; \
+            return __VA_ARGS__();                                    \
+        } else {                                                     \
+            using ActivationDtype = __nv_bfloat16;                   \
+            constexpr WeightVariant WVariant = WeightVariant::kBF16; \
+            return __VA_ARGS__();                                    \
+        }                                                            \
+    }()
+
+void moe_gemm(at::Tensor& output,
+              at::Tensor& hidden_states,
+              at::Tensor& weight,
+              c10::optional<at::Tensor>& bias,
+              at::Tensor& total_rows_before_expert,
+              int activation_raw)
+{
+    TORCH_CHECK(output.dtype() == hidden_states.dtype(),
+                "Output and hidden states must have the same dtype");
+    TORCH_CHECK(output.dtype() == weight.dtype(), "Output and weight must have the same dtype");
+
+    int64_t total_rows = hidden_states.size(0);
+    int64_t gemm_k = hidden_states.size(1);
+    int64_t gemm_n = weight.size(2);
+    int num_experts = weight.size(0);
+
+    TORCH_CHECK(total_rows == output.size(0), "Total rows dimension mismatch");
+    TORCH_CHECK(gemm_k == weight.size(1), "GEMM K dimension mismatch");
+    TORCH_CHECK(gemm_n == output.size(1), "GEMM N dimension mismatch");
+    TORCH_CHECK(num_experts == total_rows_before_expert.size(0), "Number of experts mismatch");
+
+    HIDDEN_DTYPE_SWITCH(hidden_states.dtype() == torch::kFloat16, [&] {
+        fastertransformer::MoeGemmRunner<ActivationDtype, WVariant> runner =
+            *MoeGemmContext<ActivationDtype, WVariant>::Instance().GeMM_Runner();
+
+        ActivationType activation_type = (ActivationType)activation_raw;
+        if (!bias.has_value() && activation_type == ActivationType::IDENTITY) {
+            runner.moe_gemm((ActivationDtype*)hidden_states.data_ptr(),
+                            (char*)weight.data_ptr(),
+                            nullptr,
+                            (ActivationDtype*)output.data_ptr(),
+                            (int64_t*)total_rows_before_expert.data_ptr(),
+                            total_rows,
+                            gemm_n,
+                            gemm_k,
+                            num_experts,
+                            at::cuda::getCurrentCUDAStream());
+            return;
+        } else {
+            ActivationDtype* bias_ptr = nullptr;
+            if (bias.has_value()) {
+                bias_ptr = (ActivationDtype*)bias.value().data_ptr();
+                TORCH_CHECK(num_experts == bias.value().size(0), "Number of experts mismatch");
+                TORCH_CHECK(gemm_n == bias.value().size(1), "GEMM N dimension mismatch");
+            }
+            runner.moe_gemm_bias_act((ActivationDtype*)hidden_states.data_ptr(),
+                                     (char*)weight.data_ptr(),
+                                     nullptr,
+                                     bias_ptr,
+                                     (ActivationDtype*)output.data_ptr(),
+                                     (int64_t*)total_rows_before_expert.data_ptr(),
+                                     total_rows,
+                                     gemm_n,
+                                     gemm_k,
+                                     num_experts,
+                                     activation_type,
+                                     at::cuda::getCurrentCUDAStream());
+            return;
+        }
+    });
+}
+
+#define ACT_DTYPE_SWITCH(COND, ...)                \
+    [&] {                                          \
+        if (COND) {                                \
+            using ActivationDtype = __half;        \
+            return __VA_ARGS__();                  \
+        } else {                                   \
+            using ActivationDtype = __nv_bfloat16; \
+            return __VA_ARGS__();                  \
+        }                                          \
+    }()
+
+#define WEIGHT_VARIANT_SWITCH(COND, ...)                            \
+    [&] {                                                           \
+        if (COND) {                                                 \
+            constexpr WeightVariant WVariant = WeightVariant::kFP8; \
+            return __VA_ARGS__();                                   \
+        } else {                                                    \
+            constexpr WeightVariant WVariant = WeightVariant::kFP4; \
+            return __VA_ARGS__();                                   \
+        }                                                           \
+    }()
+
+void mixed_moe_gemm(at::Tensor& output,
+                    at::Tensor& hidden_states,
+                    at::Tensor& weight,
+                    at::Tensor& scales,
+                    c10::optional<at::Tensor>& bias,
+                    at::Tensor& total_rows_before_expert,
+                    int num_bits,
+                    int activation_raw)
+{
+    TORCH_CHECK(output.dtype() == hidden_states.dtype(),
+                "Output and hidden states must have the same dtype");
+
+    int64_t total_rows = hidden_states.size(0);
+    int64_t gemm_k = hidden_states.size(1);
+    int64_t gemm_n = weight.size(2);
+    int num_experts = weight.size(0);
+
+    TORCH_CHECK(total_rows == output.size(0), "Total rows dimension mismatch");
+    TORCH_CHECK(gemm_k == weight.size(1), "GEMM K dimension mismatch");
+    TORCH_CHECK(gemm_n == output.size(1), "GEMM N dimension mismatch");
+    TORCH_CHECK(num_experts == total_rows_before_expert.size(0), "Number of experts mismatch");
+
+    ACT_DTYPE_SWITCH(hidden_states.dtype() == torch::kFloat16, [&] {
+        WEIGHT_VARIANT_SWITCH(num_bits == 8, [&] {
+            fastertransformer::MoeGemmRunner<ActivationDtype, WVariant> runner =
+                *MoeGemmContext<ActivationDtype, WVariant>::Instance().GeMM_Runner();
+
+            ActivationType activation_type = (ActivationType)activation_raw;
+            if (!bias.has_value() && activation_type == ActivationType::IDENTITY) {
+                runner.moe_gemm((ActivationDtype*)hidden_states.data_ptr(),
+                                (char*)weight.data_ptr(),
+                                (ActivationDtype*)scales.data_ptr(),
+                                (ActivationDtype*)output.data_ptr(),
+                                (int64_t*)total_rows_before_expert.data_ptr(),
+                                total_rows,
+                                gemm_n,
+                                gemm_k,
+                                num_experts,
+                                at::cuda::getCurrentCUDAStream());
+                return;
+            } else {
+                ActivationDtype* bias_ptr = nullptr;
+                if (bias.has_value()) {
+                    bias_ptr = (ActivationDtype*)bias.value().data_ptr();
+                    TORCH_CHECK(num_experts == bias.value().size(0), "Number of experts mismatch");
+                    TORCH_CHECK(gemm_n == bias.value().size(1), "GEMM N dimension mismatch");
+                }
+                runner.moe_gemm_bias_act((ActivationDtype*)hidden_states.data_ptr(),
+                                         (char*)weight.data_ptr(),
+                                         (ActivationDtype*)scales.data_ptr(),
+                                         bias_ptr,
+                                         (ActivationDtype*)output.data_ptr(),
+                                         (int64_t*)total_rows_before_expert.data_ptr(),
+                                         total_rows,
+                                         gemm_n,
+                                         gemm_k,
+                                         num_experts,
+                                         activation_type,
+                                         at::cuda::getCurrentCUDAStream());
+                return;
+            }
+        });
+    });
+}

deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm.h

@@ -0,0 +1,24 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <torch/extension.h>
+
+void moe_gemm(at::Tensor& output,
+              at::Tensor& hidden_states,
+              at::Tensor& weight,
+              c10::optional<at::Tensor>& bias,
+              at::Tensor& total_rows_before_expert,
+              int activation_raw);
+
+void mixed_moe_gemm(at::Tensor& output,
+                    at::Tensor& hidden_states,
+                    at::Tensor& weight,
+                    at::Tensor& scales,
+                    c10::optional<at::Tensor>& bias,
+                    at::Tensor& total_rows_before_expert,
+                    int num_bits,
+                    int activation_raw);

deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm.py

@@ -0,0 +1,60 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from ... import DSKernelBase
+from ....inference_utils import ActivationType, DtypeEnum
+from deepspeed.ops.op_builder import InferenceCutlassBuilder
+
+from typing import Optional
+
+
+class MoEGEMM(DSKernelBase):
+    """
+    CUTLASS implementation of MoE GEMM.
+    """
+
+    supported_dtypes = [DtypeEnum.fp16, DtypeEnum.bf16]
+    supported_act_fns = [ActivationType.GELU, ActivationType.SILU, ActivationType.RELU, ActivationType.IDENTITY]
+
+    def __init__(self, fp_dtype: DtypeEnum, act_fn: ActivationType) -> None:
+
+        if not isinstance(fp_dtype, DtypeEnum):
+            fp_dtype = DtypeEnum(fp_dtype)
+
+        if fp_dtype not in MoEGEMM.supported_dtypes:
+            raise ValueError("Unsupported data type: {}, supported_dtypes are {}".format(
+                fp_dtype, MoEGEMM.supported_dtypes))
+
+        if act_fn not in MoEGEMM.supported_act_fns:
+            raise ValueError("Unsupported activation function: {}, supported_act_fns are {}".format(
+                act_fn, MoEGEMM.supported_act_fns))
+
+        inf_module = InferenceCutlassBuilder().load()
+        self.kernel = inf_module.moe_gemm
+        self.act_fn = act_fn
+
+    def __call__(self,
+                 ordered_output: torch.Tensor,
+                 ordered_input: torch.Tensor,
+                 weights: torch.Tensor,
+                 total_rows_before_expert: torch.Tensor,
+                 biases: Optional[torch.Tensor] = None) -> None:
+        """
+            Performs a MoE GEMM. Note that the stride between token inputs must be even (the distance between byte 1 of token 0 and token 1 must be the same as the distance between byte 1 of token 1 and token 2).
+
+            Arguments:
+                ordered_output (torch.Tensor): The output of the MoE GEMM of shape [n_tokens, out_neurons].
+                ordered_input (torch.Tensor): The direct input for the MoE GEMM of shape [n_tokens, in_neurons].
+                weights (torch.Tensor): The weights of shape [n_experts, in_neurons, out_neurons]. These weights must be contiguous.
+                total_rows_before_expert (torch.Tensor): The total number of rows before each expert of shape [n_experts].
+                biases (torch.Tensor): The biases of shape [n_experts, out_neurons]. These biases must be contiguous.
+
+            Returns:
+                ordered_output
+            """
+        self.kernel(ordered_output, ordered_input, weights, biases, total_rows_before_expert, self.act_fn)
+        return ordered_output

deepspeed/inference/v2/kernels/cutlass_ops/moe_gemm/moe_gemm_api.h

@@ -0,0 +1,64 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "activation_type.h"
+#include "weight_variant.h"
+
+namespace fastertransformer {
+
+template <typename T, /*The type used for activations/scales/compute*/
+          WeightVariant V /* The type for the MoE weights */>
+class MoeGemmRunner {
+public:
+    MoeGemmRunner();
+
+    void moe_gemm_bias_act(const T* A,
+                           const char* B,
+                           const T* weight_scales,
+                           const T* biases,
+                           T* C,
+                           int64_t* total_rows_before_expert,
+                           int64_t total_rows,
+                           int64_t gemm_n,
+                           int64_t gemm_k,
+                           int num_experts,
+                           ActivationType activation_type,
+                           cudaStream_t stream);
+
+    void moe_gemm(const T* A,
+                  const char* B,
+                  const T* weight_scales,
+                  T* C,
+                  int64_t* total_rows_before_expert,
+                  int64_t total_rows,
+                  int64_t gemm_n,
+                  int64_t gemm_k,
+                  int num_experts,
+                  cudaStream_t stream);
+
+private:
+    int sm_;
+    int multi_processor_count_;
+};
+
+}  // namespace fastertransformer
+
+template <typename T, WeightVariant V>
+class MoeGemmContext {
+public:
+    MoeGemmContext() { _runner = new fastertransformer::MoeGemmRunner<T, V>(); }
+
+    virtual ~MoeGemmContext() { delete _runner; }
+
+    static MoeGemmContext& Instance()
+    {
+        static MoeGemmContext _ctx;
+        return _ctx;
+    }
+
+    fastertransformer::MoeGemmRunner<T, V>* GeMM_Runner() const { return _runner; }
+
+    fastertransformer::MoeGemmRunner<T, V>* _runner;
+};

deepspeed/inference/v2/kernels/cutlass_ops/shared_resources/weight_variant.h

@@ -0,0 +1,11 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+// Data structure that allows us to abstract internal CUTLASS datatypes/mappings
+// to the DeepSpeed-Kernels repo.
+
+#pragma once
+
+enum WeightVariant { kFP16, kBF16, kFP8, kFP4 };

deepspeed/inference/v2/kernels/ds_kernel.py

@@ -0,0 +1,32 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import ABC, abstractmethod
+
+
+class DSKernelBase(ABC):
+
+    @abstractmethod
+    def __init__(self, *args, **kwargs):
+        """
+        If necessary trigger compilation and warmup
+        Autotuning of the kernel would happen at this stage to
+        eliminate any potential hangs that might occur mid-deployment
+        Validate that the desired run configuration is compatible.
+
+        It is not necessary to call super on this method.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def __call__(self, *args, **kwargs):
+        """
+        However the kernel needs to be called, it can be called here. Auto-tuning
+        should never be performed here.
+
+        All inputs/outputs should be passed as arguments to this function. No allocations
+        should be performed here.
+        """
+        raise NotImplementedError()

deepspeed/inference/v2/kernels/includes/activation_type.h

@@ -0,0 +1,17 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+enum ActivationType {
+    GELU = 0,
+    RELU = 1,
+    SILU = 2,
+    GEGLU = 3,
+    ReGLU = 4,
+    SiGLU = 5,
+    IDENTITY = 6,
+    InvalidType = -1
+};

deepspeed/inference/v2/kernels/includes/conversion_utils.h

@@ -0,0 +1,708 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include "ds_kernel_utils.h"
+
+#include <stdint.h>
+
+#ifdef BF16_AVAILABLE
+#include <cuda_bf16.h>
+#endif
+
+namespace conversion {
+
+// Basic primitive for constructing conversions
+template <typename TO, typename FROM>
+DS_D_INLINE TO to(FROM val)
+{
+    return to(val);
+}
+
+// Specializations
+
+/********************* Identity Conversions *********************/
+/*
+Identity conversions are useful in templated functions where we might have
+a fixed destination type. For example, I might have a kernel that accepts
+__half, __nv_bfloat16, and float but always want to do the core computation
+at floating point:
+
+T mem_value = input[idx];
+float compute_value = conversion::to<float, T>(mem_value);
+
+In practice, we should be able to elide the second template parameter:
+float compute_val = conversion::to<float>(mem_value);
+
+In this case, we need an implementation to handle the T = float case
+
+NOTE: The type inferencing system appears to be unable to handle inferring the first
+template parameter, even in the trivial case.
+*/
+
+// Floating point types
+template <>
+DS_D_INLINE double to(double val)
+{
+    return val;
+}
+template <>
+DS_D_INLINE float to(float val)
+{
+    return val;
+}
+template <>
+DS_D_INLINE __half to(__half val)
+{
+    return val;
+}
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE __nv_bfloat16 to(__nv_bfloat16 val)
+{
+    return val;
+}
+#endif
+
+// Integer types
+template <>
+DS_D_INLINE int8_t to(int8_t val)
+{
+    return val;
+}
+template <>
+DS_D_INLINE uint8_t to(uint8_t val)
+{
+    return val;
+}
+template <>
+DS_D_INLINE int16_t to(int16_t val)
+{
+    return val;
+}
+template <>
+DS_D_INLINE uint16_t to(uint16_t val)
+{
+    return val;
+}
+template <>
+DS_D_INLINE int32_t to(int32_t val)
+{
+    return val;
+}
+template <>
+DS_D_INLINE uint32_t to(uint32_t val)
+{
+    return val;
+}
+template <>
+DS_D_INLINE int64_t to(int64_t val)
+{
+    return val;
+}
+template <>
+DS_D_INLINE uint64_t to(uint64_t val)
+{
+    return val;
+}
+
+// TODO: evaluate if we want bools
+
+/*********************  To Double Conversions *********************/
+
+// * to double variants
+
+// Would normally like to not use C cast, but this is an important enough conversion
+// to keep
+template <>
+DS_D_INLINE double to(float val)
+{
+#ifdef PTX_AVAILABLE
+    double ret_val;
+    asm("ctv.rn.f64.f32 %0, %1;\n" : "=d"(ret_val) : "f"(val));
+    return ret_val;
+#else
+    return double(val);
+#endif
+}
+// Note: there is a CVT instruction for __half -> double, but there's no inline interface
+// for passing a single half value
+template <>
+DS_D_INLINE double to(__half val)
+{
+    return to<double>(__half2float(val));
+}
+template <>
+DS_D_INLINE double to(int64_t val)
+{
+    return __ll2double_rn(val);
+}
+template <>
+DS_D_INLINE double to(int32_t val)
+{
+    return __int2double_rn(val);
+}
+template <>
+DS_D_INLINE double to(int16_t val)
+{
+    return __int2double_rn(val);
+}
+template <>
+DS_D_INLINE double to(int8_t val)
+{
+    return __int2double_rn(val);
+}
+template <>
+DS_D_INLINE double to(uint64_t val)
+{
+    return __ull2double_rn(val);
+}
+template <>
+DS_D_INLINE double to(uint32_t val)
+{
+    return __uint2double_rn(val);
+}
+template <>
+DS_D_INLINE double to(uint16_t val)
+{
+    return __uint2double_rn(val);
+}
+template <>
+DS_D_INLINE double to(uint8_t val)
+{
+    return __uint2double_rn(val);
+}
+
+// Same applies here
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE double to(__nv_bfloat16 val)
+{
+    return to<double>(__bfloat162float(val));
+}
+#endif
+
+/*********************  To Float Conversions *********************/
+
+template <>
+DS_D_INLINE float to(double val)
+{
+    return __double2float_rn(val);
+}
+template <>
+DS_D_INLINE float to(__half val)
+{
+    return __half2float(val);
+}
+template <>
+DS_D_INLINE float to(int64_t val)
+{
+    return __ll2float_rn(val);
+}
+template <>
+DS_D_INLINE float to(int32_t val)
+{
+    return __int2float_rn(val);
+}
+template <>
+DS_D_INLINE float to(int16_t val)
+{
+    return __int2float_rn(val);
+}
+template <>
+DS_D_INLINE float to(int8_t val)
+{
+    return __int2float_rn(val);
+}
+template <>
+DS_D_INLINE float to(uint64_t val)
+{
+    return __ull2float_rn(val);
+}
+template <>
+DS_D_INLINE float to(uint32_t val)
+{
+    return __uint2float_rn(val);
+}
+template <>
+DS_D_INLINE float to(uint16_t val)
+{
+    return __uint2float_rn(val);
+}
+template <>
+DS_D_INLINE float to(uint8_t val)
+{
+    return __uint2float_rn(val);
+}
+
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE float to(__nv_bfloat16 val)
+{
+    return __bfloat162float(val);
+}
+#endif
+
+/*********************  To Float2 Conversions *********************/
+template <>
+DS_D_INLINE float2 to(__half2 val)
+{
+    return __half22float2(val);
+}
+
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE float2 to(__nv_bfloat162 val)
+{
+    return __bfloat1622float2(val);
+}
+#endif
+
+/*********************  To Half Conversions *********************/
+template <>
+DS_D_INLINE __half to(double val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    float val_f = __double2float_rn(val);
+    return __float2half(val_f);
+#else
+    return __double2half(val);
+#endif
+}
+template <>
+DS_D_INLINE __half to(float val)
+{
+    return __float2half(val);
+}
+template <>
+DS_D_INLINE __half to(int64_t val)
+{
+    return __ll2half_rn(val);
+}
+template <>
+DS_D_INLINE __half to(int32_t val)
+{
+    return __int2half_rn(val);
+}
+template <>
+DS_D_INLINE __half to(int16_t val)
+{
+    return __short2half_rn(val);
+}
+template <>
+DS_D_INLINE __half to(int8_t val)
+{
+    return __int2half_rn(val);
+}
+template <>
+DS_D_INLINE __half to(uint64_t val)
+{
+    return __ull2half_rn(val);
+}
+template <>
+DS_D_INLINE __half to(uint32_t val)
+{
+    return __uint2half_rn(val);
+}
+template <>
+DS_D_INLINE __half to(uint16_t val)
+{
+    return __ushort2half_rn(val);
+}
+template <>
+DS_D_INLINE __half to(uint8_t val)
+{
+    return __uint2half_rn(val);
+}
+
+#ifdef BF16_AVAILABLE
+// No direct conversion
+template <>
+DS_D_INLINE __half to(__nv_bfloat16 val)
+{
+    return to<__half>(to<float>(val));
+}
+#endif
+
+/*********************  To Half2 Conversions *********************/
+template <>
+DS_D_INLINE __half2 to(float2 val)
+{
+    return __float22half2_rn(val);
+}
+template <>
+DS_D_INLINE __half2 to(float val)
+{
+    return __float2half2_rn(val);
+}
+
+#ifdef BF16_AVAILABLE
+// No direct conversion
+template <>
+DS_D_INLINE __half2 to(__nv_bfloat162 val)
+{
+    return to<__half2>(to<float2>(val));
+}
+#endif
+
+/*********************  To BF16 Conversions *********************/
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE __nv_bfloat16 to(double val)
+{
+    return __double2bfloat16(val);
+}
+template <>
+DS_D_INLINE __nv_bfloat16 to(float val)
+{
+    return __float2bfloat16(val);
+}
+template <>
+DS_D_INLINE __nv_bfloat16 to(int64_t val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __double2bfloat16(__ll2double_rn(val));
+#else
+    return __ll2bfloat16_rn(val);
+#endif
+}
+template <>
+DS_D_INLINE __nv_bfloat16 to(int32_t val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2bfloat16(__int2float_rn(val));
+#else
+    return __int2bfloat16_rn(val);
+#endif
+}
+template <>
+DS_D_INLINE __nv_bfloat16 to(int16_t val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2bfloat16(__int2float_rn(val));
+#else
+    return __short2bfloat16_rn(val);
+#endif
+}
+template <>
+DS_D_INLINE __nv_bfloat16 to(int8_t val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2bfloat16(__int2float_rn(val));
+#else
+    return __int2bfloat16_rn(val);
+#endif
+}
+template <>
+DS_D_INLINE __nv_bfloat16 to(uint64_t val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __double2bfloat16(__ull2double_rn(val));
+#else
+    return __ull2bfloat16_rn(val);
+#endif
+}
+template <>
+DS_D_INLINE __nv_bfloat16 to(uint32_t val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2bfloat16(__uint2float_rn(val));
+#else
+    return __uint2bfloat16_rn(val);
+#endif
+}
+template <>
+DS_D_INLINE __nv_bfloat16 to(uint16_t val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2bfloat16(__uint2float_rn(val));
+#else
+    return __ushort2bfloat16_rn(val);
+#endif
+}
+template <>
+DS_D_INLINE __nv_bfloat16 to(uint8_t val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2bfloat16(__uint2float_rn(val));
+#else
+    return __uint2bfloat16_rn(val);
+#endif
+}
+#endif
+
+/*********************  To BF162 Conversions *********************/
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE __nv_bfloat162 to(float2 val)
+{
+    return __float22bfloat162_rn(val);
+}
+template <>
+DS_D_INLINE __nv_bfloat162 to(float val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __bfloat162bfloat162(__float2bfloat16(val));
+#else
+    return __float2bfloat162_rn(val);
+#endif
+}
+template <>
+DS_D_INLINE __nv_bfloat162 to(__half2 val)
+{
+    return to<__nv_bfloat162>(to<float2>(val));
+}
+#endif
+
+/*********************  To INT64_T Conversions *********************/
+template <>
+DS_D_INLINE int64_t to(double val)
+{
+    return __double2ll_rn(val);
+}
+template <>
+DS_D_INLINE int64_t to(float val)
+{
+    return __float2ll_rn(val);
+}
+template <>
+DS_D_INLINE int64_t to(__half val)
+{
+    return __half2ll_rn(val);
+}
+// No direct support for integer casts at the C++ level and I don't feel they're so important
+// to demand an PTX at this time
+
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE int64_t to(__nv_bfloat16 val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2ll_rn(__bfloat162float(val));
+#else
+    return __bfloat162ll_rn(val);
+#endif
+}
+#endif
+
+/*********************  To INT32_T Conversions *********************/
+template <>
+DS_D_INLINE int32_t to(double val)
+{
+    return __double2int_rn(val);
+}
+template <>
+DS_D_INLINE int32_t to(float val)
+{
+    return __float2int_rn(val);
+}
+template <>
+DS_D_INLINE int32_t to(__half val)
+{
+    return __half2int_rn(val);
+}
+// No direct support for integer casts at the C++ level and I don't feel they're so important
+// to demand an PTX at this time
+
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE int32_t to(__nv_bfloat16 val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2int_rn(__bfloat162float(val));
+#else
+    return __bfloat162int_rn(val);
+#endif
+}
+#endif
+
+/*********************  To INT16_T Conversions *********************/
+template <>
+DS_D_INLINE int16_t to(double val)
+{
+    return __double2int_rn(val);
+}
+template <>
+DS_D_INLINE int16_t to(float val)
+{
+    return __float2int_rn(val);
+}
+template <>
+DS_D_INLINE int16_t to(__half val)
+{
+    return __half2int_rn(val);
+}
+// No direct support for integer casts at the C++ level and I don't feel they're so important
+// to demand an PTX at this time
+
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE int16_t to(__nv_bfloat16 val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2int_rn(__bfloat162float(val));
+#else
+    return __bfloat162int_rn(val);
+#endif
+}
+#endif
+
+/*********************  To INT8_T Conversions *********************/
+template <>
+DS_D_INLINE int8_t to(double val)
+{
+    return __double2int_rn(val);
+}
+template <>
+DS_D_INLINE int8_t to(float val)
+{
+    return __float2int_rn(val);
+}
+template <>
+DS_D_INLINE int8_t to(__half val)
+{
+    return __half2int_rn(val);
+}
+// No direct support for integer casts at the C++ level and I don't feel they're so important
+// to demand an PTX at this time
+
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE int8_t to(__nv_bfloat16 val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2int_rn(__bfloat162float(val));
+#else
+    return __bfloat162int_rn(val);
+#endif
+}
+#endif
+
+/*********************  To UINT64_T Conversions *********************/
+template <>
+DS_D_INLINE uint64_t to(double val)
+{
+    return __double2ull_rn(val);
+}
+template <>
+DS_D_INLINE uint64_t to(float val)
+{
+    return __float2ull_rn(val);
+}
+template <>
+DS_D_INLINE uint64_t to(__half val)
+{
+    return __half2ull_rn(val);
+}
+// No direct support for integer casts at the C++ level and I don't feel they're so important
+// to demand an PTX at this time
+
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE uint64_t to(__nv_bfloat16 val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2ull_rn(__bfloat162float(val));
+#else
+    return __bfloat162ull_rn(val);
+#endif
+}
+#endif
+
+/*********************  To UINT32_T Conversions *********************/
+template <>
+DS_D_INLINE uint32_t to(double val)
+{
+    return __double2uint_rn(val);
+}
+template <>
+DS_D_INLINE uint32_t to(float val)
+{
+    return __float2uint_rn(val);
+}
+template <>
+DS_D_INLINE uint32_t to(__half val)
+{
+    return __half2uint_rn(val);
+}
+// No direct support for integer casts at the C++ level and I don't feel they're so important
+// to demand an PTX at this time
+
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE uint32_t to(__nv_bfloat16 val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2uint_rn(__bfloat162float(val));
+#else
+    return __bfloat162uint_rn(val);
+#endif
+}
+#endif
+
+/*********************  To UINT16_T Conversions *********************/
+template <>
+DS_D_INLINE uint16_t to(double val)
+{
+    return __double2uint_rn(val);
+}
+template <>
+DS_D_INLINE uint16_t to(float val)
+{
+    return __float2uint_rn(val);
+}
+template <>
+DS_D_INLINE uint16_t to(__half val)
+{
+    return __half2uint_rn(val);
+}
+// No direct support for integer casts at the C++ level and I don't feel they're so important
+// to demand an PTX at this time
+
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE uint16_t to(__nv_bfloat16 val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2uint_rn(__bfloat162float(val));
+#else
+    return __bfloat162uint_rn(val);
+#endif
+}
+#endif
+
+/*********************  To UINT8_T Conversions *********************/
+template <>
+DS_D_INLINE uint8_t to(double val)
+{
+    return __double2uint_rn(val);
+}
+template <>
+DS_D_INLINE uint8_t to(float val)
+{
+    return __float2uint_rn(val);
+}
+template <>
+DS_D_INLINE uint8_t to(__half val)
+{
+    return __half2uint_rn(val);
+}
+// No direct support for integer casts at the C++ level and I don't feel they're so important
+// to demand an PTX at this time
+
+#ifdef BF16_AVAILABLE
+template <>
+DS_D_INLINE uint8_t to(__nv_bfloat16 val)
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __float2uint_rn(__bfloat162float(val));
+#else
+    return __bfloat162uint_rn(val);
+#endif
+}
+#endif
+
+}  // namespace conversion

deepspeed/inference/v2/kernels/includes/ds_kernel_utils.h

@@ -0,0 +1,58 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/*
+Centralized header file for preprocessor macros and constants
+used throughout the codebase.
+*/
+
+#pragma once
+
+#include <cuda.h>
+#include <cuda_fp16.h>
+
+#ifdef BF16_AVAILABLE
+#include <cuda_bf16.h>
+#endif
+
+#define DS_HD_INLINE __host__ __device__ __forceinline__
+#define DS_D_INLINE __device__ __forceinline__
+
+#ifdef __HIP_PLATFORM_AMD__
+
+// constexpr variant of warpSize for templating
+constexpr int hw_warp_size = ROCM_WAVEFRONT_SIZE;
+#define HALF_PRECISION_AVAILABLE = 1
+#include <hip/hip_cooperative_groups.h>
+#include <hip/hip_fp16.h>
+
+#else  // !__HIP_PLATFORM_AMD__
+
+// constexpr variant of warpSize for templating
+constexpr int hw_warp_size = 32;
+
+#if __CUDA_ARCH__ >= 530
+#define HALF_PRECISION_AVAILABLE = 1
+#define PTX_AVAILABLE
+#endif  // __CUDA_ARCH__ >= 530
+
+#if __CUDA_ARCH__ >= 800
+#define ASYNC_COPY_AVAILABLE
+#endif  // __CUDA_ARCH__ >= 800
+
+#include <cooperative_groups.h>
+#include <cuda_fp16.h>
+
+#endif  //__HIP_PLATFORM_AMD__
+
+inline int next_pow2(const int val)
+{
+    int rounded_val = val - 1;
+    rounded_val |= rounded_val >> 1;
+    rounded_val |= rounded_val >> 2;
+    rounded_val |= rounded_val >> 4;
+    rounded_val |= rounded_val >> 8;
+    return rounded_val + 1;
+}

deepspeed/inference/v2/kernels/includes/reduction_utils.h

@@ -0,0 +1,778 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include "conversion_utils.h"
+#include "ds_kernel_utils.h"
+#include "memory_access_utils.h"
+
+namespace cg = cooperative_groups;
+
+namespace reduce {
+
+enum class ROpType {
+    // Addition
+    Add,
+
+    // Maximum reduction
+    Max,
+
+    // Minimum reduction
+    Min,
+};
+
+constexpr int max_threads = 1024;
+constexpr int max_warps = max_threads / hw_warp_size;
+
+/*
+High level API. The API takes in a set of operations and variables
+and performs that reduction operation on that variable. The reductions
+of each of the arguments are completely independent of each other (
+i.e., the val1-op1 combination has no impact on val2-op2).
+
+Example usage:
+``` cpp
+float max_val;
+float min_val;
+reduce::block<rop::Max, rop::Min>(tb, warp, max_val, min_val);
+```
+
+TODO(cmikeh2): In theory, we might be able to do this sequentially with
+device functions and rely on the assembler correctly behaving. My initial
+instinct is this won't work, but if it does it would reduce implementation
+cost significantly.
+
+TODO(cmikeh2): We need to support sub-block reductions. The warp intrinsic
+currently supports this (more incidentally than anything else). It is not
+uncommon in something like softmax or a fused attention kernel to map multiple
+reductions to a thread block, but each reduction itself is only scoped
+to part of the threads (i.e block size = 512, 128 threads per reduction).
+*/
+template <ROpType Op, int warp_bound = max_warps>
+DS_D_INLINE void block(cg::thread_block& tb, cg::thread_block_tile<hw_warp_size>& warp, float& val);
+
+template <ROpType Op1, ROpType Op2, int warp_bound = max_warps>
+DS_D_INLINE void block(cg::thread_block& tb,
+                       cg::thread_block_tile<hw_warp_size>& warp,
+                       float& val1,
+                       float& val2);
+
+template <ROpType Op1, ROpType Op2, ROpType Op3, int warp_bound = max_warps>
+DS_D_INLINE void block(cg::thread_block& tb,
+                       cg::thread_block_tile<hw_warp_size>& warp,
+                       float& val1,
+                       float& val2,
+                       float& val3);
+
+template <ROpType Op1, ROpType Op2, ROpType Op3, ROpType Op4, int warp_bound = max_warps>
+DS_D_INLINE void block(cg::thread_block& tb,
+                       cg::thread_block_tile<hw_warp_size>& warp,
+                       float& val1,
+                       float& val2,
+                       float& val3,
+                       float& val4);
+
+/*
+The partitioned block is a special case of the above where in the warps of a threadblock are
+partitioned into separate independent reductions. For example, I might have an 8 warp thread block
+in which each pair of warps is processing an independent piece of data. I would then reduce that
+data with the something like the following:
+``` cpp
+float max_val;
+reduce::partitioned_block<rop::Max, 2>(tb, warp, max_val);
+```
+After which, each pair of warps would have coherent data with each other. Note, this API will not
+provide correct results if the number of warps per partition is not a power of 2.
+*/
+template <ROpType Op, int num_threads>
+DS_D_INLINE void partitioned_block(cg::thread_block& tb,
+                                   cg::thread_block_tile<hw_warp_size>& warp,
+                                   float& val);
+
+template <ROpType Op1, ROpType Op2, int num_threads>
+DS_D_INLINE void partitioned_block(cg::thread_block& tb,
+                                   cg::thread_block_tile<hw_warp_size>& warp,
+                                   float& val1,
+                                   float& val2);
+
+template <ROpType Op1, ROpType Op2, ROpType Op3, int num_threads>
+DS_D_INLINE void partitioned_block(cg::thread_block& tb,
+                                   cg::thread_block_tile<hw_warp_size>& warp,
+                                   float& val1,
+                                   float& val2,
+                                   float& val3);
+
+template <ROpType Op1, ROpType Op2, ROpType Op3, ROpType Op4, int num_threads>
+DS_D_INLINE void partitioned_block(cg::thread_block& tb,
+                                   cg::thread_block_tile<hw_warp_size>& warp,
+                                   float& val1,
+                                   float& val2,
+                                   float& val3,
+                                   float& val4);
+
+/*
+Single element reduction primitives. Used inside serial collection
+loops.
+
+Example usage:
+using rop = reduce::OpType;
+float min = init<rop::Min>();
+for (int i = 0; i < 4; i++) {
+    min = reduce::element<rop::Min>(min, data[i]);
+}
+*/
+
+template <ROpType Op, typename T>
+DS_D_INLINE T element(const T lhs, const T rhs);
+
+template <ROpType OType, typename T = float>
+DS_D_INLINE T init();
+
+/********************** Internal reduction APIs **********************/
+
+/*
+Single element "reductions". TODO(cmikeh2): this sort of "op" concept
+should be refactored into its own implementation at some point. This interface
+may be easily expanded for new types/operations, but the typical reductions
+we need are covered with min/max/add on float.
+
+NOTE: there is no mean reduction because that relies on knowledge of how
+many values were already reduced into each scalar. Implementing this on top
+of reduce should be straightforward (can just wrap the sum reduction) and
+would be a good extension of the header.
+*/
+
+DS_D_INLINE int _warp_rank()
+{
+    const int thread_rank =
+        threadIdx.x + threadIdx.y * blockDim.x + threadIdx.z * blockDim.x * blockDim.y;
+    return thread_rank / hw_warp_size;
+}
+
+/* Float element reduce implementations */
+template <>
+DS_D_INLINE float element<ROpType::Add>(const float lhs, const float rhs)
+{
+    return lhs + rhs;
+}
+
+template <>
+DS_D_INLINE float element<ROpType::Max>(const float lhs, const float rhs)
+{
+    return fmaxf(lhs, rhs);
+}
+
+template <>
+DS_D_INLINE float element<ROpType::Min>(const float lhs, const float rhs)
+{
+    return fminf(lhs, rhs);
+}
+
+/* __half element reduce implementation */
+template <>
+DS_D_INLINE __half element<ROpType::Add>(const __half lhs, const __half rhs)
+{
+    return lhs + rhs;
+}
+
+template <>
+DS_D_INLINE __half element<ROpType::Max>(const __half lhs, const __half rhs)
+{
+#if __CUDA_ARCH__ >= 800
+    // Intrinsic limited to Ampere + newer
+    return __hmax(lhs, rhs);
+#else
+    return (lhs > rhs) ? lhs : rhs;
+#endif
+}
+
+template <>
+DS_D_INLINE __half element<ROpType::Min>(const __half lhs, const __half rhs)
+{
+#if __CUDA_ARCH__ >= 800
+    // Intrinsic limited to Ampere + newer
+    return __hmin(lhs, rhs);
+#else
+    return (lhs < rhs) ? lhs : rhs;
+#endif
+}
+
+/* __half2 element reduce implementation */
+template <>
+DS_D_INLINE __half2 element<ROpType::Add>(const __half2 lhs, const __half2 rhs)
+{
+    return lhs + rhs;
+}
+
+template <>
+DS_D_INLINE __half2 element<ROpType::Max>(const __half2 lhs, const __half2 rhs)
+{
+#if __CUDA_ARCH__ >= 800
+    return __hmax2(lhs, rhs);
+#else
+    __half2 ret_val;
+    ret_val.x = (lhs.x > rhs.x) ? lhs.x : rhs.x;
+    ret_val.y = (lhs.y > rhs.y) ? lhs.y : rhs.y;
+    return ret_val;
+#endif
+}
+
+template <>
+DS_D_INLINE __half2 element<ROpType::Min>(const __half2 lhs, const __half2 rhs)
+{
+#if __CUDA_ARCH__ >= 800
+    return __hmin2(lhs, rhs);
+#else
+    __half2 ret_val;
+    ret_val.x = (lhs.x < rhs.x) ? lhs.x : rhs.x;
+    ret_val.y = (lhs.y < rhs.y) ? lhs.y : rhs.y;
+    return ret_val;
+#endif
+}
+
+template <>
+DS_D_INLINE int32_t element<ROpType::Add>(const int32_t lhs, const int32_t rhs)
+{
+    return lhs + rhs;
+}
+
+template <>
+DS_D_INLINE int32_t element<ROpType::Max>(const int32_t lhs, const int32_t rhs)
+{
+    return (lhs > rhs) ? lhs : rhs;
+}
+
+template <>
+DS_D_INLINE int32_t element<ROpType::Min>(const int32_t lhs, const int32_t rhs)
+{
+    return (lhs < rhs) ? lhs : rhs;
+}
+
+template <>
+DS_D_INLINE uint32_t element<ROpType::Add>(const uint32_t lhs, const uint32_t rhs)
+{
+    return lhs + rhs;
+}
+
+template <>
+DS_D_INLINE uint32_t element<ROpType::Max>(const uint32_t lhs, const uint32_t rhs)
+{
+    return (lhs > rhs) ? lhs : rhs;
+}
+
+template <>
+DS_D_INLINE uint32_t element<ROpType::Min>(const uint32_t lhs, const uint32_t rhs)
+{
+    return (lhs < rhs) ? lhs : rhs;
+}
+
+template <>
+DS_D_INLINE int64_t element<ROpType::Add>(const int64_t lhs, const int64_t rhs)
+{
+    return lhs + rhs;
+}
+
+template <>
+DS_D_INLINE int64_t element<ROpType::Max>(const int64_t lhs, const int64_t rhs)
+{
+    return (lhs > rhs) ? lhs : rhs;
+}
+
+template <>
+DS_D_INLINE int64_t element<ROpType::Min>(const int64_t lhs, const int64_t rhs)
+{
+    return (lhs < rhs) ? lhs : rhs;
+}
+
+/*
+Reduction initialization primitives
+*/
+template <>
+DS_D_INLINE float init<ROpType::Add>()
+{
+    return 0.0f;
+}
+
+template <>
+DS_D_INLINE float init<ROpType::Min>()
+{
+    // Positive infinity
+    return INFINITY;
+}
+
+template <>
+DS_D_INLINE float init<ROpType::Max>()
+{
+    // Negative infinity
+    return -INFINITY;
+}
+
+template <>
+DS_D_INLINE __half init<ROpType::Add>()
+{
+    constexpr __half_raw zero = {0x0000};
+    return __half(zero);
+}
+
+template <>
+DS_D_INLINE __half init<ROpType::Min>()
+{
+    constexpr __half_raw inf = {0x7C00};
+    return __half(inf);
+}
+
+template <>
+DS_D_INLINE __half init<ROpType::Max>()
+{
+    constexpr __half_raw neg_inf = {0xFC00};
+    return __half(neg_inf);
+}
+
+template <>
+DS_D_INLINE __half2 init<ROpType::Add>()
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __half2{_Float16_2{0x0000, 0x0000}};
+#else
+    constexpr __half2_raw zero = {0x0000, 0x0000};
+    return __half2(zero);
+#endif
+}
+
+template <>
+DS_D_INLINE __half2 init<ROpType::Min>()
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __half2{_Float16_2{0x7C00, 0x7C00}};
+#else
+    constexpr __half2_raw inf = {0x7C00, 0x7C00};
+    return __half2(inf);
+#endif
+}
+
+template <>
+DS_D_INLINE __half2 init<ROpType::Max>()
+{
+#ifdef __HIP_PLATFORM_AMD__
+    return __half2{_Float16_2{0xFC00, 0xFC00}};
+#else
+    constexpr __half2_raw neg_inf = {0xFC00, 0xFC00};
+    return __half2(neg_inf);
+#endif
+}
+
+template <>
+DS_D_INLINE int32_t init<ROpType::Add>()
+{
+    return 0;
+}
+
+template <>
+DS_D_INLINE int32_t init<ROpType::Min>()
+{
+    return 0x7FFFFFFF;
+}
+
+template <>
+DS_D_INLINE int32_t init<ROpType::Max>()
+{
+    return 0x80000000;
+}
+
+template <>
+DS_D_INLINE uint32_t init<ROpType::Add>()
+{
+    return 0;
+}
+
+template <>
+DS_D_INLINE uint32_t init<ROpType::Min>()
+{
+    return 0xFFFFFFFF;
+}
+
+template <>
+DS_D_INLINE uint32_t init<ROpType::Max>()
+{
+    return 0;
+}
+
+template <>
+DS_D_INLINE int64_t init<ROpType::Add>()
+{
+    return 0;
+}
+
+template <>
+DS_D_INLINE int64_t init<ROpType::Min>()
+{
+    return 0x7FFFFFFFFFFFFFFF;
+}
+
+template <>
+DS_D_INLINE int64_t init<ROpType::Max>()
+{
+    return 0x8000000000000000;
+}
+
+template <>
+DS_D_INLINE uint64_t init<ROpType::Add>()
+{
+    return 0;
+}
+
+template <>
+DS_D_INLINE uint64_t init<ROpType::Min>()
+{
+    return 0xFFFFFFFFFFFFFFFF;
+}
+
+template <>
+DS_D_INLINE uint64_t init<ROpType::Max>()
+{
+    return 0;
+}
+
+template <ROpType Op, typename T>
+DS_D_INLINE void init(T* data)
+{
+    data[0] = init<Op, T>();
+}
+
+template <ROpType Op1, ROpType Op2, typename T>
+DS_D_INLINE void init(T* data)
+{
+    data[0] = init<Op1, T>();
+    data[1] = init<Op2, T>();
+}
+
+template <ROpType Op1, ROpType Op2, ROpType Op3, typename T>
+DS_D_INLINE void init(T* data)
+{
+    data[0] = init<Op1, T>();
+    data[1] = init<Op2, T>();
+    data[2] = init<Op3, T>();
+}
+
+template <ROpType Op1, ROpType Op2, ROpType Op3, ROpType Op4, typename T>
+DS_D_INLINE void init(T* data)
+{
+    data[0] = init<Op1, T>();
+    data[1] = init<Op2, T>();
+    data[2] = init<Op3, T>();
+    data[3] = init<Op4, T>();
+}
+
+/*
+Warp reduction primitives
+
+`reduction_width` is an unsafe template parameter, that is that
+when using `reduction_width` < hw_warp_size the warp is partitioned
+into `hw_warp_size` / `reduction_width` groups of partial sums.
+
+If someone can figure out how to use variadic templates in a reasonable way
+here (fold is C++17 only and I don't think helps and recursion feels like
+huge overkill that harms readability) that would be wonderful.
+*/
+
+template <typename T, ROpType Op, int reduce_width = hw_warp_size>
+DS_D_INLINE void _warp(cg::thread_block_tile<hw_warp_size>& warp, T* data)
+{
+#pragma unroll
+    for (int i = 1; i < reduce_width; i *= 2) {
+        data[0] = element<Op>(data[0], warp.shfl_xor(data[0], i));
+    }
+}
+
+template <typename T, ROpType Op1, ROpType Op2, int reduce_width = hw_warp_size>
+DS_D_INLINE void _warp(cg::thread_block_tile<hw_warp_size>& warp, T* data)
+{
+#pragma unroll
+    for (int i = 1; i < reduce_width; i *= 2) {
+        data[0] = element<Op1>(data[0], warp.shfl_xor(data[0], i));
+        data[1] = element<Op2>(data[1], warp.shfl_xor(data[1], i));
+    }
+}
+
+template <typename T, ROpType Op1, ROpType Op2, ROpType Op3, int reduce_width = hw_warp_size>
+DS_D_INLINE void _warp(cg::thread_block_tile<hw_warp_size>& warp, T* data)
+{
+#pragma unroll
+    for (int i = 1; i < reduce_width; i *= 2) {
+        data[0] = element<Op1>(data[0], warp.shfl_xor(data[0], i));
+        data[1] = element<Op2>(data[1], warp.shfl_xor(data[1], i));
+        data[2] = element<Op3>(data[2], warp.shfl_xor(data[2], i));
+    }
+}
+
+template <typename T,
+          ROpType Op1,
+          ROpType Op2,
+          ROpType Op3,
+          ROpType Op4,
+          int reduce_width = hw_warp_size>
+DS_D_INLINE void _warp(cg::thread_block_tile<hw_warp_size>& warp, T* data)
+{
+#pragma unroll
+    for (int i = 1; i < reduce_width; i *= 2) {
+        data[0] = element<Op1>(data[0], warp.shfl_xor(data[0], i));
+        data[1] = element<Op2>(data[1], warp.shfl_xor(data[1], i));
+        data[2] = element<Op3>(data[2], warp.shfl_xor(data[2], i));
+        data[3] = element<Op4>(data[3], warp.shfl_xor(data[3], i));
+    }
+}
+
+/*
+Implementation for primary block reduction that serves both `block` and
+`partitioned_block`.
+
+Total warps refers to the reduction width of the reduction, not
+the number of warps in the block (which may exceed that
+if the block is partitioned or if we do a conservative bound at
+compile time).
+*/
+template <typename T, int total_warps, ROpType... Ops>
+DS_D_INLINE void _block(cg::thread_block& tb,
+                        cg::thread_block_tile<hw_warp_size>& warp_arg,
+                        T* data)
+{
+    constexpr int elems = sizeof...(Ops);
+    constexpr int bytes = sizeof(T);
+    // Unused when `partition_size == 1` or total_warps == 1
+    __shared__ T reduce_buffer[max_warps * elems];
+
+#ifdef __HIP_PLATFORM_AMD__
+    const int total_threads = blockDim.x * blockDim.y * blockDim.z;
+    const int running_warps = total_threads / hw_warp_size;
+#else
+    const int running_warps = warp_arg.meta_group_size();
+#endif
+
+    // Always perform warp-scope reduction
+    _warp<T, Ops...>(warp_arg, data);
+
+    // If max_warps == 1 let's skip the runtime check
+    if (total_warps != 1) {
+        if (warp_arg.thread_rank() == 0) {
+#pragma unroll
+            for (int i = 0; i < elems; i++) {
+                mem_access::store_shared<bytes>(reduce_buffer + elems * _warp_rank() + i, data + i);
+            }
+        }
+
+        // Synchronization inside block-uniform conditional is safe
+        tb.sync();
+
+        if (_warp_rank() == 0) {
+            if (warp_arg.thread_rank() < running_warps) {
+#pragma unroll
+                for (int i = 0; i < elems; i++) {
+                    mem_access::load_shared<bytes>(
+                        data + i, reduce_buffer + elems * warp_arg.thread_rank() + i);
+                }
+            } else {
+                init<Ops...>(data);
+            }
+
+            _warp<T, Ops..., total_warps>(warp_arg, data);
+
+#pragma unroll
+            for (int i = 0; i < elems; i++) {
+                mem_access::store_shared<bytes>(reduce_buffer + elems * warp_arg.thread_rank() + i,
+                                                data + i);
+            }
+        }
+
+        // Synchronization inside block-uniform conditional is safe
+        tb.sync();
+
+#pragma unroll
+        for (int i = 0; i < elems; i++) {
+            mem_access::load_shared<bytes>(data + i, reduce_buffer + _warp_rank() * elems + i);
+        }
+    }
+}
+
+/*
+Main API implementations. For the most part, they just convert the individual
+variables into arrays, which makes working with them easier with a single
+implementation. In theory, we could use the `_block` implementation as another
+option, but the nature of using a pointer is a little less safe and this allows
+us to obfuscate the details of the partitioned implementation.
+*/
+template <ROpType Op, int warp_bound>
+DS_D_INLINE void block(cg::thread_block& tb, cg::thread_block_tile<hw_warp_size>& warp, float& val)
+{
+    _block<float, warp_bound, Op>(tb, warp, &val);
+}
+
+template <ROpType Op1, ROpType Op2, int warp_bound>
+DS_D_INLINE void block(cg::thread_block& tb,
+                       cg::thread_block_tile<hw_warp_size>& warp,
+                       float& val1,
+                       float& val2)
+{
+    float data[2] = {val1, val2};
+    _block<float, warp_bound, Op1, Op2>(tb, warp, data);
+    val1 = data[0];
+    val2 = data[1];
+}
+
+template <ROpType Op1, ROpType Op2, ROpType Op3, int warp_bound>
+DS_D_INLINE void block(cg::thread_block& tb,
+                       cg::thread_block_tile<hw_warp_size>& warp,
+                       float& val1,
+                       float& val2,
+                       float& val3)
+{
+    float data[3] = {val1, val2, val3};
+    _block<float, warp_bound, Op1, Op2, Op3>(tb, warp, data);
+    val1 = data[0];
+    val2 = data[1];
+    val3 = data[2];
+}
+
+template <ROpType Op1, ROpType Op2, ROpType Op3, ROpType Op4, int warp_bound>
+DS_D_INLINE void block(cg::thread_block& tb,
+                       cg::thread_block_tile<hw_warp_size>& warp,
+                       float& val1,
+                       float& val2,
+                       float& val3,
+                       float& val4)
+{
+    float data[4] = {val1, val2, val3, val4};
+    _block<float, warp_bound, Op1, Op2, Op3, Op4>(tb, warp, data);
+    val1 = data[0];
+    val2 = data[1];
+    val3 = data[2];
+    val4 = data[3];
+}
+
+/*
+Note: for the partitioned blocks, the implementation does not support non-power of 2 blocks in order
+to shorten block scale reduction length.
+*/
+template <ROpType Op, int num_threads>
+DS_D_INLINE void partitioned_block(cg::thread_block& tb,
+                                   cg::thread_block_tile<hw_warp_size>& warp,
+                                   float& val)
+{
+    if (num_threads <= hw_warp_size) {
+        _warp<float, Op, num_threads>(warp, &val);
+    } else {
+        constexpr int num_warps = num_threads / hw_warp_size;
+        _block<float, num_warps, Op>(tb, warp, &val);
+    }
+}
+
+template <ROpType Op1, ROpType Op2, int num_threads>
+DS_D_INLINE void partitioned_block(cg::thread_block& tb,
+                                   cg::thread_block_tile<hw_warp_size>& warp,
+                                   float& val1,
+                                   float& val2)
+{
+    float data[2] = {val1, val2};
+
+    if (num_threads <= hw_warp_size) {
+        _warp<float, Op1, Op2, num_threads>(warp, data);
+    } else {
+        constexpr int num_warps = num_threads / hw_warp_size;
+        _block<float, num_warps, Op1, Op2>(tb, warp, data);
+    }
+
+    val1 = data[0];
+    val2 = data[1];
+}
+
+template <ROpType Op1, ROpType Op2, ROpType Op3, int num_threads>
+DS_D_INLINE void partitioned_block(cg::thread_block& tb,
+                                   cg::thread_block_tile<hw_warp_size>& warp,
+                                   float& val1,
+                                   float& val2,
+                                   float& val3)
+{
+    float data[3] = {val1, val2, val3};
+
+    if (num_threads <= hw_warp_size) {
+        _warp<float, Op1, Op2, Op3, num_threads>(warp, data);
+    } else {
+        constexpr int num_warps = num_threads / hw_warp_size;
+        _block<float, num_warps, Op1, Op2, Op3>(tb, warp, data);
+    }
+
+    val1 = data[0];
+    val2 = data[1];
+    val3 = data[2];
+}
+
+template <ROpType Op1, ROpType Op2, ROpType Op3, ROpType Op4, int num_threads>
+DS_D_INLINE void partitioned_block(cg::thread_block& tb,
+                                   cg::thread_block_tile<hw_warp_size>& warp,
+                                   float& val1,
+                                   float& val2,
+                                   float& val3,
+                                   float& val4)
+{
+    float data[4] = {val1, val2, val3, val4};
+
+    if (num_threads <= hw_warp_size) {
+        _warp<float, Op1, Op2, Op3, Op4, num_threads>(warp, data);
+    } else {
+        constexpr int num_warps = num_threads / hw_warp_size;
+        _block<float, num_warps, Op1, Op2, Op3, Op4>(tb, warp, data);
+    }
+
+    val1 = data[0];
+    val2 = data[1];
+    val3 = data[2];
+    val4 = data[3];
+}
+
+/*
+Arg-reduce is a specialization of the above. We only support this with a single reduction
+parameter. This only works for max/min reductions.
+*/
+
+__align__(8) struct IdxReduceResult {
+    /*
+    NOTE: ORDERING MATTERS HERE! The idx is the least significant set of bits
+    and the val is the most significant. Changing the order of this declaration
+    will break the code.
+    */
+    int idx;
+    float val;
+};
+
+template <ROpType Op, int warpBound>
+DS_D_INLINE IdxReduceResult
+idx_reduce(cg::thread_block& tb, cg::thread_block_tile<hw_warp_size>& warp, float val, int idx)
+{
+    IdxReduceResult res = {idx, val};
+
+    // Clear out the nan. This shouldn't be an issue for our initial applications
+    if (isnan(val)) res.val = init<Op>();
+
+    // Can do float compares as integers. By packing the index into the lower bits
+    // we can just do a single int64 rather than a branch, compare, and select.
+    // One side benefit of this is that it is by nature a stable algorithm and
+    // will always bias ties to the higher index.
+    int64_t* res_as_int = reinterpret_cast<int64_t*>(&res);
+
+    // The way floating point compare works is normally to perform a sign comparison
+    // and if they match, then do a comparison of the rest of the bits as unsigned
+    // integers. Since we are bundling these, that means for negative values we need
+    // to reverse the sort order, which we can do with an XOR.
+    if (val < 0) { *res_as_int ^= 0x7fffffff00000000; }
+
+    _block<int64_t, warpBound, Op>(tb, warp, res_as_int);
+
+    // Sign bit is preserved, so we can check if we need to invert the mantissa back
+    if (res.val < 0) { *res_as_int ^= 0x7fffffff00000000; }
+
+    return res;
+}
+
+}  // namespace reduce

deepspeed/inference/v2/kernels/ragged_ops/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .atom_builder import *
+from .blocked_flash import *
+from .embed import *
+from .linear_blocked_kv_rotary import *
+from .logits_gather import *
+from .moe_gather import *
+from .moe_scatter import *
+from .top_k_gating import *

deepspeed/inference/v2/kernels/ragged_ops/atom_builder/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .atom_builder import *

deepspeed/inference/v2/kernels/ragged_ops/atom_builder/atom_builder.cpp

@@ -0,0 +1,53 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "atom_builder.h"
+#include "attention_atom.h"
+#include "ragged_dtypes.h"
+
+int32_t build_atoms(torch::Tensor& atoms_ten,
+                    torch::Tensor& batch_metadata,
+                    torch::Tensor& seq_metadata,
+                    torch::Tensor& kv_ptrs,
+                    const int32_t q_block_size,
+                    const int32_t kv_block_size)
+{
+    const RaggedBatchDescriptor* batch_desc =
+        reinterpret_cast<const RaggedBatchDescriptor*>(batch_metadata.data_ptr());
+
+    const InflightSeqDescriptor* seq_desc =
+        reinterpret_cast<const InflightSeqDescriptor*>(seq_metadata.data_ptr());
+
+    int32_t** kv_ptr_list = reinterpret_cast<int32_t**>(kv_ptrs.data_ptr());
+
+    AttentionAtom* atoms = reinterpret_cast<AttentionAtom*>(atoms_ten.data_ptr());
+
+    int32_t n_atoms = 0;
+    for (int i = 0; i < batch_desc->n_sequences; i++) {
+        const int seq_atoms = (seq_desc[i].n_tokens + q_block_size - 1) / q_block_size;
+        int32_t cur_start_idx = seq_desc[i].start_idx;
+        int32_t global_start_idx = seq_desc[i].seen_tokens;
+        int32_t remaining_toks = seq_desc[i].n_tokens;
+
+        for (int j = 0; j < seq_atoms; j++) {
+            atoms[n_atoms].block_idx_list = kv_ptr_list[i];
+            atoms[n_atoms].q_start_idx = cur_start_idx;
+            atoms[n_atoms].q_len = std::min(remaining_toks, q_block_size);
+            atoms[n_atoms].global_q_idx = global_start_idx;
+
+            const int32_t end_toks = global_start_idx + atoms[n_atoms].q_len;
+            // TODO(cmikeh2): This logic needs to be changed for sparse implementations
+            atoms[n_atoms].kv_blocks = (end_toks + kv_block_size - 1) / kv_block_size;
+            atoms[n_atoms].total_extent = end_toks;
+
+            cur_start_idx += atoms[n_atoms].q_len;
+            global_start_idx += atoms[n_atoms].q_len;
+            remaining_toks -= atoms[n_atoms].q_len;
+            n_atoms++;
+        }
+    }
+
+    return n_atoms;
+}

deepspeed/inference/v2/kernels/ragged_ops/atom_builder/atom_builder.h

@@ -0,0 +1,21 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <torch/extension.h>
+
+/*
+Construct the attention atoms given the ragged metadata for the current batch.
+This could largely be done at the Python level, but since we pack the KV ptr
+alongside the int32_t metadata, it gets very ugly to handle the mixed-width
+data structures (since we're packing them in a single tensor).
+*/
+int32_t build_atoms(torch::Tensor& atoms_ten,
+                    torch::Tensor& batch_metadata,
+                    torch::Tensor& seq_metadata,
+                    torch::Tensor& kv_ptrs,
+                    const int32_t q_block_size,
+                    const int32_t kv_block_size);

deepspeed/inference/v2/kernels/ragged_ops/atom_builder/atom_builder.py

@@ -0,0 +1,50 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Tuple
+
+import torch
+
+from ... import DSKernelBase
+from deepspeed.ops.op_builder import RaggedOpsBuilder
+from ....ragged import RaggedBatchWrapper
+
+
+class AtomBuilder(DSKernelBase):
+    """
+    C++ implementation to populate the attention atoms for the blocked attention
+    kernel.
+    """
+
+    def __init__(self) -> None:
+        """
+        Triggers compilation of the C++ implementation.
+        """
+        inf_module = RaggedOpsBuilder().load()
+        self.kernel = inf_module.build_atoms
+
+    def __call__(self, atoms: torch.Tensor, ragged_batch: RaggedBatchWrapper, q_block_size: int,
+                 kv_block_size: int) -> Tuple[torch.Tensor, int]:
+        """
+        Populates the attention atoms for the blocked attention kernel.
+
+        Args:
+            atoms (torch.Tensor): Pre-allocated int32 tensor of shape [max_atoms, 8]
+            ragged_batch (torch.Tensor): Wrapper for the ragged batch.
+            q_block_size (int): The block size for the queries (as determined by the
+                attention implementation)
+            kv_block_size (int): The block size for the keys/values (as determined by the
+                attention implementation)
+
+        Returns:
+
+        """
+        if atoms.device != torch.device("cpu"):
+            raise RuntimeError("AtomBuilder must be called on tensors")
+
+        n_atoms = self.kernel(atoms, ragged_batch.batch_metadata_buffer(on_device=False),
+                              ragged_batch.inflight_seq_descriptors(on_device=False),
+                              ragged_batch.kv_ptrs(on_device=False), q_block_size, kv_block_size)
+        return atoms, n_atoms

deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .blocked_flash import *

deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/attention_atom.h

@@ -0,0 +1,39 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <cstdint>
+#include "cuda.h"
+
+struct AttentionAtom {
+    /*
+    The attention atom describes the workload of a particular query. The attention
+    kernel will execute each ``AttentionAtom`` for each head of the model.
+    */
+
+    // Pointer to a list of KV block indices.
+    int32_t* block_idx_list;
+
+    // Index of first token in the ragged batch associated with this atom.
+    int32_t q_start_idx;
+
+    // Number of tokens in the ragged batch associated with this atom.
+    int32_t q_len;
+
+    // Number of key/value blocks associated with this atom. All but the last are
+    // assumed to be fully dense.
+    int32_t kv_blocks;
+
+    // Number of tokens in the last key/value block.
+    int32_t total_extent;
+
+    // Global index of the first token in the atom. For example, in a prompt continuation
+    // in which we have already processed 768 tokens, this would be 768.
+    int32_t global_q_idx;
+
+    // Unused
+    int32_t unused;
+};

deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.cpp

@@ -0,0 +1,101 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/extension.h>
+
+#include "blocked_flash.h"
+#include "flash.h"
+
+#define CHECK_SHAPE(x, ...)                                     \
+    TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), \
+                #x " must have shape (" #__VA_ARGS__ ")")
+
+void flash_attn_by_atoms(at::Tensor& out,
+                         at::Tensor& q,
+                         at::Tensor& k,
+                         at::Tensor& v,
+                         at::Tensor& attention_atoms,
+                         const float softmax_scale,
+                         const bool is_causal)
+{
+    auto dprops = at::cuda::getCurrentDeviceProperties();
+
+    bool is_sm8x = dprops->major == 8 && dprops->minor >= 0;
+    bool is_sm90 = dprops->major == 9 && dprops->minor == 0;
+    TORCH_CHECK(is_sm90 || is_sm8x, "FlashAttention only supports Ampere GPUs or newer.");
+
+    auto q_dtype = q.dtype();
+    TORCH_CHECK(q_dtype == torch::kFloat16 || q_dtype == torch::kBFloat16,
+                "FlashAttention only support fp16 and bf16 data type");
+    if (q_dtype == torch::kBFloat16) {
+        TORCH_CHECK(is_sm90 || is_sm8x, "bfloat16 is only supported on Ampere GPUs or newer");
+    }
+    TORCH_CHECK(k.dtype() == q_dtype, "query and key must have the same dtype");
+    TORCH_CHECK(v.dtype() == q_dtype, "query and value must have the same dtype");
+
+    TORCH_CHECK(q.is_cuda(), "Input tensor must be on CUDA device");
+    TORCH_CHECK(k.is_cuda(), "Input tensor must be on CUDA device");
+    TORCH_CHECK(v.is_cuda(), "Input tensor must be on CUDA device");
+
+    TORCH_CHECK(q.stride(-1) == 1, "Input tensor must have contiguous last dimension");
+    TORCH_CHECK(k.stride(-1) == 1, "Input tensor must have contiguous last dimension");
+    TORCH_CHECK(v.stride(-1) == 1, "Input tensor must have contiguous last dimension");
+
+    const int total_q = q.size(0);
+    const int head_size = k.size(-1);
+    const int num_heads_kv = k.size(-2);
+    const int num_heads_q = q.size(-1) / head_size;
+
+    TORCH_CHECK(head_size <= 256, "head_size must be <= 256");
+    TORCH_CHECK(head_size % 8 == 0, "head_size must be divisible by 8");
+    TORCH_CHECK(num_heads_q % num_heads_kv == 0, "num_heads_q must be divisible by num_heads_kv");
+
+    Flash_fwd_params params;
+
+    params.is_bf16 = q.dtype() == torch::kBFloat16;
+
+    // Set the pointers and strides.
+    params.q_ptr = q.data_ptr();
+    params.k_ptr = k.data_ptr();
+    params.v_ptr = v.data_ptr();
+    params.o_ptr = out.data_ptr();
+    params.atoms = reinterpret_cast<AttentionAtom*>(attention_atoms.data_ptr());
+
+    // All stride are in elements, not bytes.
+    params.q_row_stride = q.stride(0);
+    params.k_row_stride = k.stride(1);
+    params.v_row_stride = v.stride(1);
+    params.o_row_stride = out.stride(0);
+
+    // Assume heads are contiguous.
+    params.q_head_stride = head_size;
+    params.k_head_stride = head_size;
+    params.v_head_stride = head_size;
+    params.o_head_stride = head_size;
+
+    // Head params
+    params.h = num_heads_q;
+    params.h_k = num_heads_kv;
+    params.h_h_k_ratio = num_heads_q / num_heads_kv;
+    params.d = head_size;
+    auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; };
+    params.d_rounded = round_multiple(head_size, 32);
+    params.num_atoms = attention_atoms.size(0);
+
+    // Set the different scale values.
+    params.scale_softmax = softmax_scale;
+    params.scale_softmax_log2 = softmax_scale * M_LOG2E;
+
+    params.is_causal = is_causal;
+
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    run_mha_fwd(params, stream);
+}

deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.h

@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <torch/extension.h>
+
+void flash_attn_by_atoms(at::Tensor& out,
+                         at::Tensor& q,
+                         at::Tensor& k,
+                         at::Tensor& v,
+                         at::Tensor& attention_atoms,
+                         const float softmax_scale,
+                         const bool is_causal);

deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/blocked_flash.py

@@ -0,0 +1,107 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from ....inference_utils import DtypeEnum
+from deepspeed.ops.op_builder import RaggedOpsBuilder
+
+from ... import DSKernelBase
+
+
+def get_q_block_size(head_size: int) -> int:
+    """
+    Returns the query block size required by the kernel given a head size.
+    """
+    cc_major, cc_minor = torch.cuda.get_device_capability(get_accelerator().current_device())  #ignore-cuda
+
+    if cc_major < 8:
+        raise RuntimeError("Blocked attention requires CUDA compute capability >= 8.0")
+
+    if head_size <= 64:
+        return 128
+    elif head_size <= 160:
+        if cc_minor != 0:
+            return 64
+        else:
+            return 128
+    elif head_size == 192:
+        return 128
+    elif head_size == 224:
+        if cc_minor != 0:
+            return 64
+        else:
+            return 128
+    else:
+        if cc_major == 8 and cc_minor == 0:
+            return 128
+        else:
+            return 64
+
+
+def get_kv_block_size(head_size: int) -> int:
+    """
+    Return preferred granulatity for blocked KV-cache implementation.
+    """
+    cc_major, cc_minor = torch.cuda.get_device_capability(get_accelerator().current_device())  #ignore-cuda
+
+    if cc_major < 8:
+        raise RuntimeError("Blocked attention requires CUDA compute capability >= 8.0")
+
+    if head_size <= 64:
+        return 128
+    elif head_size != 160 or cc_minor != 0:
+        return 64
+    else:
+        return 32
+
+
+class BlockedFlashAttn(DSKernelBase):
+    """
+    Modified implementation of flash-attn-2 tuned for inference on blocked KV-cache and wider
+    range of input sequence lengths.
+    """
+
+    supported_dtypes = [DtypeEnum.fp16, DtypeEnum.bf16]
+
+    def __init__(self, head_size: int, dtype: DtypeEnum) -> None:
+        """
+        Triggers any compilation of the kernels.
+        """
+        if not isinstance(dtype, DtypeEnum):
+            dtype = DtypeEnum(dtype)
+
+        if dtype not in BlockedFlashAttn.supported_dtypes:
+            raise ValueError("Unsupported data type: {}, supported data types are {}".format(
+                dtype, BlockedFlashAttn.supported_dtypes))
+
+        # For testing, need to revert to 32
+        if head_size % 16 != 0:
+            raise ValueError("Head size must be divisible by 32 (configured with {})".format(head_size))
+
+        inf_module = RaggedOpsBuilder().load()
+        self.kernel = inf_module.flash_attn_by_atoms
+
+    def __call__(self, out: torch.Tensor, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, atoms: torch.Tensor,
+                 softmax_scale: float) -> torch.Tensor:
+        """
+        Flash attention implementation atop a blocked KV-cache. Atoms should be pre-populated.
+        See attention_atom.h for further details on the structure of the information.
+
+        Arguments:
+            out (torch.Tensor): Output tensor of shape [tokens, hidden_size]
+            q (torch.Tensor): Query tensor of shape [tokens, hidden_size]
+            k (torch.Tensor): Key cache tensor of shape [n_blocks, block_size, n_heads_kv, head_size]. This Tensor only needs to be contiguous on the final dimension.
+            v (torch.Tensor): Value cache tensor of shape [n_blocks, block_size, n_heads_kv, head_size]. This Tensor only needs to be contiguous on the final dimension.
+            atoms (torch.Tensor): Atom information tensor of shape [num_atoms, 8] and type int32.
+                Not all data is readable in this format. See attention_atom.h for further details.
+            softmax_scale (float): Softmax scale factor.
+
+        Returns:
+            out (torch.Tensor): Output tensor of shape [tokens, hidden_size]
+        """
+        self.kernel(out, q, k, v, atoms, softmax_scale, True)
+        return out

deepspeed/inference/v2/kernels/ragged_ops/blocked_flash/flash.h

@@ -0,0 +1,74 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+/******************************************************************************
+Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <cuda.h>
+#include <vector>
+
+#include "attention_atom.h"
+
+constexpr int TOTAL_DIM = 0;
+constexpr int H_DIM = 1;
+constexpr int D_DIM = 2;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Qkv_params {
+    using index_t = uint32_t;
+    // The QKV matrices.
+    void* __restrict__ q_ptr;
+    void* __restrict__ k_ptr;
+    void* __restrict__ v_ptr;
+
+    // The stride between rows of the Q, K and V matrices.
+    index_t q_row_stride;
+    index_t k_row_stride;
+    index_t v_row_stride;
+    index_t q_head_stride;
+    index_t k_head_stride;
+    index_t v_head_stride;
+
+    // The number of heads.
+    int h, h_k;
+    // In the case of multi-query and grouped-query attention (MQA/GQA), nheads_k could be
+    // different from nheads (query).
+    int h_h_k_ratio;  // precompute h / h_k,
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Flash_fwd_params : public Qkv_params {
+    // The O matrix (output).
+    void* __restrict__ o_ptr;
+
+    // The attention metadata
+    AttentionAtom* __restrict__ atoms;
+
+    // Total attention atoms
+    int num_atoms;
+
+    // The stride between rows of O.
+    index_t o_row_stride;
+    index_t o_head_stride;
+
+    // The dimensions
+    int d, d_rounded;
+
+    // The scaling factors for the kernel.
+    float scale_softmax;
+    float scale_softmax_log2;
+
+    bool is_bf16;
+    bool is_causal;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void run_mha_fwd(Flash_fwd_params& params, cudaStream_t stream);

deepspeed/inference/v2/kernels/ragged_ops/embed/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .embed import RaggedEmbeddingKernel

deepspeed/inference/v2/kernels/ragged_ops/embed/embed.cpp

@@ -0,0 +1,101 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "embed.h"
+#include "ragged_kernel_helpers.h"
+
+#ifdef BF16_AVAILABLE
+#define DISPATCH_FOR_FLOAT(DTYPE, ...)                       \
+    [&] {                                                    \
+        if (DTYPE == torch::kFloat32) {                      \
+            using float_t = float;                           \
+            return __VA_ARGS__();                            \
+        } else if (DTYPE == torch::kFloat16) {               \
+            using float_t = __half;                          \
+            return __VA_ARGS__();                            \
+        } else if (DTYPE == torch::kBFloat16) {              \
+            using float_t = __nv_bfloat16;                   \
+            return __VA_ARGS__();                            \
+        } else {                                             \
+            TORCH_CHECK(false, "Unsupported dispatch type"); \
+        }                                                    \
+    }()
+#else
+#define DISPATCH_FOR_FLOAT(DTYPE, ...)                       \
+    [&] {                                                    \
+        if (DTYPE == torch::kFloat32) {                      \
+            using float_t = float;                           \
+            return __VA_ARGS__();                            \
+        } else if (DTYPE == torch::kFloat16) {               \
+            using float_t = __half;                          \
+            return __VA_ARGS__();                            \
+        } else {                                             \
+            TORCH_CHECK(false, "Unsupported dispatch type"); \
+        }                                                    \
+    }()
+#endif
+
+#define DISPATCH_FOR_INT(DTYPE, ...)                         \
+    [&] {                                                    \
+        if (DTYPE == torch::kInt32) {                        \
+            using int_t = int32_t;                           \
+            return __VA_ARGS__();                            \
+        } else if (DTYPE == torch::kInt64) {                 \
+            using int_t = int64_t;                           \
+            return __VA_ARGS__();                            \
+        } else {                                             \
+            TORCH_CHECK(false, "Unsupported dispatch type"); \
+        }                                                    \
+    }()
+
+/*
+Embeddings kernel aware of ragged batch structure.
+*/
+void ragged_embed(torch::Tensor& embedded_tokens,
+                  torch::Tensor& input_ids,
+                  torch::Tensor& embedding_weight,
+                  c10::optional<torch::Tensor>& position_embedding_weight,
+                  int32_t pos_embed_offset,
+                  torch::Tensor& batch_metadata,
+                  torch::Tensor& seq_metadata,
+                  torch::Tensor& tokens_to_seq,
+                  torch::Tensor& kv_ptrs)
+{
+    // We don't care about KV cache here, so just hardcoding 0s for block_size/num_blocks
+    BatchWrapperCPP batch_wrapper =
+        make_cpp_batch_wrapper(batch_metadata, seq_metadata, tokens_to_seq, kv_ptrs, 0, 0);
+
+    const int32_t n_tokens = input_ids.numel();
+    const int32_t embed_dim = embedding_weight.size(1);
+    const int32_t vocab_size = embedding_weight.size(0);
+
+    DISPATCH_FOR_INT(input_ids.scalar_type(), [&] {
+        DISPATCH_FOR_FLOAT(embedding_weight.scalar_type(), [&] {
+            float_t* pos_embed_ptr = nullptr;
+            int32_t max_position_embed_idx = 0;
+            if (position_embedding_weight.has_value()) {
+                TORCH_CHECK(
+                    position_embedding_weight.value().options().dtype() ==
+                        embedding_weight.options().dtype(),
+                    "position_embedding_weight and embedding_weight must have the same dtype");
+                pos_embed_ptr =
+                    reinterpret_cast<float_t*>(position_embedding_weight.value().data_ptr());
+                max_position_embed_idx = position_embedding_weight.value().size(0) - 1;
+            }
+
+            launch_ragged_embed_kernel((float_t*)embedded_tokens.data_ptr(),
+                                       (const int_t*)input_ids.data_ptr(),
+                                       (const float_t*)embedding_weight.data_ptr(),
+                                       pos_embed_ptr,
+                                       batch_wrapper,
+                                       n_tokens,
+                                       embed_dim,
+                                       vocab_size,
+                                       max_position_embed_idx,
+                                       pos_embed_offset,
+                                       at::cuda::getCurrentCUDAStream());
+        });
+    });
+}

deepspeed/inference/v2/kernels/ragged_ops/embed/embed.h

@@ -0,0 +1,23 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+#include "embed.cuh"
+
+/*
+Embeddings kernel aware of ragged batch structure.
+*/
+void ragged_embed(torch::Tensor& embedded_tokens,
+                  torch::Tensor& input_ids,
+                  torch::Tensor& embedding_weight,
+                  c10::optional<torch::Tensor>& position_weight,
+                  int32_t position_embed_offset,
+                  torch::Tensor& batch_metadata,
+                  torch::Tensor& seq_metadata,
+                  torch::Tensor& tokens_to_seq,
+                  torch::Tensor& kv_ptrs);

deepspeed/inference/v2/kernels/ragged_ops/embed/embed.py

@@ -0,0 +1,67 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional
+
+import torch
+
+from ... import DSKernelBase
+from deepspeed.ops.op_builder import RaggedOpsBuilder
+from ....inference_utils import elem_size
+from ....ragged import RaggedBatchWrapper
+
+
+class RaggedEmbeddingKernel(DSKernelBase):
+    """
+    Ragged-aware CUDA kernel implementation for an embedding lookup. This will only lookup
+    the necessary tokens for a padded batch (i.e. if we are CGed and running with a slightly
+    larger batch size than the actual tokens).
+    """
+
+    supported_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+    supported_token_dtypes = [torch.int32, torch.int64]
+
+    def __init__(self, embed_dtype: torch.dtype, token_dtype: torch.dtype, embed_dim: int) -> None:
+        """
+        Args:
+            fp_dtype (torch.dtype): Data type of the embedding table and output dtype.
+                Supported values are torch.float16, torch.bfloat16, and torch.float32.
+            token_dtype (torch.dtype): Data type of the token ids. Supported values are
+                torch.int32 and torch.int64.
+            embed_dim (int): Embedding dimension. Must be aligned to 16 bytes.
+        """
+        if embed_dtype not in RaggedEmbeddingKernel.supported_dtypes:
+            raise ValueError("Unsupported embedding data type: {}, supported_dtypes are {}".format(
+                embed_dtype, RaggedEmbeddingKernel.supported_dtypes))
+
+        if token_dtype not in RaggedEmbeddingKernel.supported_token_dtypes:
+            raise ValueError("Unsupported token data type: {}, supported_dtypes are {}".format(
+                token_dtype, RaggedEmbeddingKernel.supported_token_dtypes))
+
+        if elem_size(embed_dtype) * embed_dim % 16 != 0:
+            raise ValueError("Embedding dimension must be aligned to 16 bytes, got {}".format(embed_dim))
+
+        inf_module = RaggedOpsBuilder().load()
+        self.kernel = inf_module.ragged_embed
+
+    def __call__(self,
+                 embedded_tokens: torch.Tensor,
+                 ragged_wrapper: RaggedBatchWrapper,
+                 embedding_weight: torch.Tensor,
+                 position_embed_weight: Optional[torch.Tensor] = None,
+                 position_embed_offset: int = 0) -> torch.Tensor:
+        """
+        Ragged aware embedding lookup.
+
+        Args:
+            embedded_tokens (torch.Tensor): Output tensor of shape [num_tokens, embed_dim]
+            ragged_wrapper (RaggedBatchWrapper): Wrapper for the ragged batch.
+            embedding_weight (torch.Tensor): Embedding table of shape [vocab_size, embed_dim]
+        """
+        self.kernel(embedded_tokens, ragged_wrapper.input_ids(),
+                    embedding_weight, position_embed_weight, position_embed_offset,
+                    ragged_wrapper.batch_metadata_buffer(), ragged_wrapper.inflight_seq_descriptors(),
+                    ragged_wrapper.tokens_to_seq(), ragged_wrapper.kv_ptrs())
+        return embedded_tokens

deepspeed/inference/v2/kernels/ragged_ops/embed/embed_cuda.cu

@@ -0,0 +1,137 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "ds_kernel_utils.h"
+#include "embed.cuh"
+#include "memory_access_utils.h"
+#include "ragged_dtypes.h"
+
+namespace embed {
+
+constexpr int granularity = 16;
+constexpr int threads = 512;
+
+}  // namespace embed
+
+template <typename TokenType, typename EmbedType>
+__global__ void ragged_embed_kernel(EmbedType* embedded_tokens,
+                                    const TokenType* input_ids,
+                                    const EmbedType* embedding_weight,
+                                    const EmbedType* position_weight,
+                                    const BatchWrapperCPP batch_desc,
+                                    const int32_t embed_dim,
+                                    const int32_t vocab_size,
+                                    const int32_t max_position_embed_idx,
+                                    const int32_t position_embed_offset)
+{
+    constexpr int T_vector = embed::granularity / sizeof(EmbedType);
+
+    const int32_t token_idx = blockIdx.y;
+
+    // It's possible our batch is padded (under CG conditions typically)
+    if (token_idx >= batch_desc.batch_metadata->n_tokens) return;
+
+    TokenType token_value = input_ids[token_idx];
+
+    if (token_value >= vocab_size || token_value < 0) {
+        // TODO(cmikeh2): This is invalid, but not sure how we want to handle it being invalid
+        // yet.
+        return;
+    }
+
+    const EmbedType* embedding_row = embedding_weight + token_value * embed_dim;
+    EmbedType* dest_row = embedded_tokens + token_idx * embed_dim;
+
+    const int channel_offset = (threadIdx.x + embed::threads * blockIdx.x) * T_vector;
+
+    if (channel_offset < embed_dim) {
+        EmbedType reg_buf[T_vector];
+
+        mem_access::load_global<embed::granularity>(reg_buf, embedding_row + channel_offset);
+
+        if (position_weight != nullptr) {
+            // Map the token to its global idx (indirect memory accesses aren't great but whatever)
+            const int32_t seq_idx = batch_desc.tokens_to_seq[token_idx];
+            const InflightSeqDescriptor seq_desc = batch_desc.seq_metadata[seq_idx];
+            int32_t pos_emb_idx = seq_desc.seen_tokens + (token_idx - seq_desc.start_idx);
+
+            // Position embed offset is an OPT-specific feature I think?
+            pos_emb_idx = pos_emb_idx + position_embed_offset;
+
+            // This clamping is technically
+            pos_emb_idx = (pos_emb_idx < 0) ? 0 : pos_emb_idx;
+            pos_emb_idx = (pos_emb_idx >= max_position_embed_idx) ? max_position_embed_idx
+                                                                  : pos_emb_idx;
+
+            const EmbedType* position_embedding_row = position_weight + pos_emb_idx * embed_dim;
+
+            EmbedType pos_buf[T_vector];
+            mem_access::load_global<embed::granularity>(pos_buf,
+                                                        position_embedding_row + channel_offset);
+
+#pragma unroll
+            for (int i = 0; i < T_vector; i++) { reg_buf[i] += pos_buf[i]; }
+        }
+
+        mem_access::store_global<embed::granularity>(dest_row + channel_offset, reg_buf);
+    }
+}
+
+template <typename TokenType, typename EmbedType>
+void launch_ragged_embed_kernel(EmbedType* embedded_tokens,
+                                const TokenType* input_ids,
+                                const EmbedType* embedding_weight,
+                                const EmbedType* position_weight,
+                                const BatchWrapperCPP batch_desc,
+                                const int32_t n_tokens,
+                                const int32_t embed_dim,
+                                const int32_t vocab_size,
+                                const int32_t max_position_embed_idx,
+                                const int32_t position_embed_offset,
+                                cudaStream_t stream)
+{
+    constexpr int T_vector = embed::granularity / sizeof(EmbedType);
+    constexpr int elems_per_block = embed::threads * T_vector;
+    const int parallel_blocks = (embed_dim + elems_per_block - 1) / elems_per_block;
+
+    const dim3 grid_dim(parallel_blocks, n_tokens, 1);
+    const dim3 block_dim(embed::threads, 1, 1);
+
+    ragged_embed_kernel<TokenType, EmbedType>
+        <<<grid_dim, block_dim, 0, stream>>>(embedded_tokens,
+                                             input_ids,
+                                             embedding_weight,
+                                             position_weight,
+                                             batch_desc,
+                                             embed_dim,
+                                             vocab_size,
+                                             max_position_embed_idx,
+                                             position_embed_offset);
+}
+
+#define INSTANTIATE_EMBED_FOR_TYPES(TOKEN_TYPE, EMBED_TYPE)           \
+    template void launch_ragged_embed_kernel<TOKEN_TYPE, EMBED_TYPE>( \
+        EMBED_TYPE * embedded_tokens,                                 \
+        const TOKEN_TYPE* input_ids,                                  \
+        const EMBED_TYPE* embedding_weight,                           \
+        const EMBED_TYPE* position_weight,                            \
+        const BatchWrapperCPP batch_descriptor,                       \
+        const int32_t n_tokens,                                       \
+        const int32_t embed_dim,                                      \
+        const int32_t vocab_size,                                     \
+        const int32_t max_position_embed_idx,                         \
+        const int32_t position_embed_offset,                          \
+        cudaStream_t stream);
+
+INSTANTIATE_EMBED_FOR_TYPES(int32_t, float)
+INSTANTIATE_EMBED_FOR_TYPES(int64_t, float)
+
+INSTANTIATE_EMBED_FOR_TYPES(int32_t, __half)
+INSTANTIATE_EMBED_FOR_TYPES(int64_t, __half)
+
+#ifdef BF16_AVAILABLE
+INSTANTIATE_EMBED_FOR_TYPES(int32_t, __nv_bfloat16)
+INSTANTIATE_EMBED_FOR_TYPES(int64_t, __nv_bfloat16)
+#endif

deepspeed/inference/v2/kernels/ragged_ops/includes/top_k_utils.h

@@ -0,0 +1,21 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#define TOP_K_SWITCH(N_TOP_K, ...)         \
+    [&] {                                  \
+        if (1 == N_TOP_K) {                \
+            constexpr int CONST_TOP_K = 1; \
+            __VA_ARGS__();                 \
+        } else if (2 == N_TOP_K) {         \
+            constexpr int CONST_TOP_K = 2; \
+            __VA_ARGS__();                 \
+        } else if (4 == N_TOP_K) {         \
+            constexpr int CONST_TOP_K = 4; \
+            __VA_ARGS__();                 \
+        } else if (8 == N_TOP_K) {         \
+            constexpr int CONST_TOP_K = 8; \
+            __VA_ARGS__();                 \
+        }                                  \
+    }()

deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .blocked_kv_rotary import *
+from .blocked_trained_kv_rotary import *
+from .linear_blocked_kv_copy import *

deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.cpp

@@ -0,0 +1,195 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#include "blocked_kv_rotary.h"
+#include "ragged_kernel_helpers.h"
+
+#define DISPATCH_KV_ROTARY(T_TYPE, C_TYPE)                                 \
+    if (q.options().dtype() == torch::T_TYPE) {                            \
+        launch_kv_rotary_kernel<C_TYPE>((C_TYPE*)kv_cache.data_ptr(),      \
+                                        (C_TYPE*)q.data_ptr(),             \
+                                        (C_TYPE*)k.data_ptr(),             \
+                                        (C_TYPE*)v.data_ptr(),             \
+                                        (C_TYPE*)inv_freq_ptr,             \
+                                        rotary_dim,                        \
+                                        theta_base,                        \
+                                        batch_wrapper,                     \
+                                        qkv_stride,                        \
+                                        kv_cache_stride,                   \
+                                        v_offset,                          \
+                                        inv_freq_stride,                   \
+                                        q_ratio,                           \
+                                        head_size,                         \
+                                        n_tokens,                          \
+                                        n_q_heads,                         \
+                                        at::cuda::getCurrentCUDAStream()); \
+    }
+
+/*
+Rotary position embeddings + copy into KV cache. This implementation assumes
+that the inverse frequencies should be ready from global memory rather than
+synthesized in the kernel.
+
+Arguments:
+    kv_cache: [n_blocks, block_size, 2, n_kv_heads, head_size]
+    q: [n_tokens, n_q_heads * head_size]
+    k: [n_tokens, n_kv_heads * head_size]
+    v: [n_tokens, n_kv_heads * head_size]
+    inv_freq: [max_seq_len, head_size // 2]
+*/
+void kv_trained_rotary_embeddings(torch::Tensor& kv_cache,
+                                  torch::Tensor& q,
+                                  torch::Tensor& k,
+                                  torch::Tensor& v,
+                                  torch::Tensor& inv_freq,
+                                  torch::Tensor& batch_metadata,
+                                  torch::Tensor& seq_metadata,
+                                  torch::Tensor& tokens_to_seq,
+                                  torch::Tensor& kv_ptrs)
+{
+    const int32_t n_tokens = q.size(0);
+    TORCH_CHECK(n_tokens == k.size(0));
+    TORCH_CHECK(n_tokens == v.size(0));
+
+    const float theta_base = 0.f;
+    const int32_t rotary_dim = inv_freq.size(0) * 2;
+
+    // Dimensions
+    const int32_t block_size = kv_cache.size(1);
+    const int32_t n_kv_heads = kv_cache.size(3);
+    const int32_t head_size = kv_cache.size(4);
+
+    // Strides
+    const int32_t qkv_stride = q.stride(0);              // Per token
+    const int32_t kv_cache_stride = kv_cache.stride(1);  // Per token
+    const int32_t v_offset = kv_cache.stride(2);         // From k_cache to v_cache
+    const int32_t inv_freq_stride = inv_freq.stride(0);  // Per token idx
+
+    const int n_q_heads = q.size(1) / head_size;
+    const int q_ratio = n_q_heads / n_kv_heads;
+
+    void* inv_freq_ptr = (void*)inv_freq.data_ptr();
+
+    BatchWrapperCPP batch_wrapper = make_cpp_batch_wrapper(
+        batch_metadata, seq_metadata, tokens_to_seq, kv_ptrs, block_size, kv_cache.size(0));
+
+    DISPATCH_KV_ROTARY(kHalf, __half);
+
+#ifdef BF16_AVAILABLE
+    DISPATCH_KV_ROTARY(kBFloat16, __nv_bfloat16);
+#endif
+}
+
+/*
+Rotary position embeddings + copy into KV cache. This implementation assumes
+that the inverse frequencies should be synthesized in the kernel.
+
+Arguments:
+    kv_cache: [n_blocks, block_size, 2, n_kv_heads, head_size]
+    q: [n_tokens, n_q_heads * head_size]
+    k: [n_tokens, n_kv_heads * head_size]
+    v: [n_tokens, n_kv_heads * head_size]
+*/
+void kv_rotary_embeddings(torch::Tensor& kv_cache,
+                          torch::Tensor& q,
+                          torch::Tensor& k,
+                          torch::Tensor& v,
+                          const int32_t rotary_dim,
+                          const float theta_base,
+                          torch::Tensor& batch_metadata,
+                          torch::Tensor& seq_metadata,
+                          torch::Tensor& tokens_to_seq,
+                          torch::Tensor& kv_ptrs)
+{
+    const int32_t n_tokens = q.size(0);
+    TORCH_CHECK(n_tokens == k.size(0));
+    TORCH_CHECK(n_tokens == v.size(0));
+
+    // Dimensions
+    const int32_t block_size = kv_cache.size(1);
+    const int32_t n_kv_heads = kv_cache.size(3);
+    const int32_t head_size = kv_cache.size(4);
+
+    // Strides
+    const int32_t qkv_stride = q.stride(0);              // Per token
+    const int32_t kv_cache_stride = kv_cache.stride(1);  // Per token
+    const int32_t v_offset = kv_cache.stride(2);         // From k_cache to v_cache
+    const int32_t inv_freq_stride = 0;                   // Per token idx
+
+    const int n_q_heads = q.size(1) / head_size;
+    const int q_ratio = n_q_heads / n_kv_heads;
+
+    void* inv_freq_ptr = nullptr;
+
+    BatchWrapperCPP batch_wrapper = make_cpp_batch_wrapper(
+        batch_metadata, seq_metadata, tokens_to_seq, kv_ptrs, block_size, kv_cache.size(0));
+
+    DISPATCH_KV_ROTARY(kHalf, __half);
+
+#ifdef BF16_AVAILABLE
+    DISPATCH_KV_ROTARY(kBFloat16, __nv_bfloat16);
+#endif
+}
+
+#define DISPATCH_KV_COPY(T_TYPE, C_TYPE)                                 \
+    if (q.options().dtype() == torch::T_TYPE) {                          \
+        launch_kv_copy_kernel<C_TYPE>((C_TYPE*)kv_cache.data_ptr(),      \
+                                      (C_TYPE*)q.data_ptr(),             \
+                                      (C_TYPE*)k.data_ptr(),             \
+                                      (C_TYPE*)v.data_ptr(),             \
+                                      batch_wrapper,                     \
+                                      qkv_stride,                        \
+                                      kv_cache_stride,                   \
+                                      v_offset,                          \
+                                      q_ratio,                           \
+                                      head_size,                         \
+                                      n_tokens,                          \
+                                      n_q_heads,                         \
+                                      at::cuda::getCurrentCUDAStream()); \
+    }
+
+/*
+Copy into linear KV cache.
+*/
+void linear_kv_copy(torch::Tensor& kv_cache,
+                    torch::Tensor& q,
+                    torch::Tensor& k,
+                    torch::Tensor& v,
+                    torch::Tensor& batch_metadata,
+                    torch::Tensor& seq_metadata,
+                    torch::Tensor& tokens_to_seq,
+                    torch::Tensor& kv_ptrs)
+{
+    const int32_t n_tokens = q.size(0);
+    TORCH_CHECK(n_tokens == k.size(0));
+    TORCH_CHECK(n_tokens == v.size(0));
+
+    // Dimensions
+    const int32_t block_size = kv_cache.size(1);
+    const int32_t n_kv_heads = kv_cache.size(3);
+    const int32_t head_size = kv_cache.size(4);
+
+    // Strides
+    const int32_t qkv_stride = q.stride(0);  // Per token
+    TORCH_CHECK(qkv_stride == k.stride(0));
+    TORCH_CHECK(qkv_stride == v.stride(0));
+
+    const int32_t kv_cache_stride = kv_cache.stride(1);  // Per token
+    const int32_t v_offset = kv_cache.stride(2);         // From k_cache to v_cache
+
+    const int n_q_heads = q.size(1) / head_size;
+
+    TORCH_CHECK(n_q_heads % n_kv_heads == 0);
+    const int q_ratio = n_q_heads / n_kv_heads;
+
+    BatchWrapperCPP batch_wrapper = make_cpp_batch_wrapper(
+        batch_metadata, seq_metadata, tokens_to_seq, kv_ptrs, block_size, kv_cache.size(0));
+
+    DISPATCH_KV_COPY(kHalf, __half);
+
+#ifdef BF16_AVAILABLE
+    DISPATCH_KV_COPY(kBFloat16, __nv_bfloat16);
+#endif
+}

deepspeed/inference/v2/kernels/ragged_ops/linear_blocked_kv_rotary/blocked_kv_rotary.h

@@ -0,0 +1,65 @@
+// Copyright (c) Microsoft Corporation.
+// SPDX-License-Identifier: Apache-2.0
+
+// DeepSpeed Team
+
+#pragma once
+
+#include <c10/cuda/CUDAStream.h>
+#include <torch/extension.h>
+#include "blocked_kv_rotary.cuh"
+
+/*
+Rotary position embeddings + copy into KV cache. This implementation assumes
+that the inverse frequencies should be ready from global memory rather than
+synthesized in the kernel.
+
+Arguments:
+    kv_cache: [n_blocks, block_size, 2, n_kv_heads, head_size]
+    q: [n_tokens, n_q_heads * head_size]
+    k: [n_tokens, n_kv_heads * head_size]
+    v: [n_tokens, n_kv_heads * head_size]
+    inv_freq: [max_seq_len, head_size // 2]
+*/
+void kv_trained_rotary_embeddings(torch::Tensor& kv_cache,
+                                  torch::Tensor& q,
+                                  torch::Tensor& k,
+                                  torch::Tensor& v,
+                                  torch::Tensor& inv_freq,
+                                  torch::Tensor& batch_metadata,
+                                  torch::Tensor& seq_metadata,
+                                  torch::Tensor& tokens_to_seq,
+                                  torch::Tensor& kv_ptrs);
+
+/*
+Rotary position embeddings + copy into KV cache. This implementation assumes
+that the inverse frequencies should be synthesized in the kernel.
+
+Arguments:
+    kv_cache: [n_blocks, block_size, 2, n_kv_heads, head_size]
+    q: [n_tokens, n_q_heads * head_size]
+    k: [n_tokens, n_kv_heads * head_size]
+    v: [n_tokens, n_kv_heads * head_size]
+*/
+void kv_rotary_embeddings(torch::Tensor& kv_cache,
+                          torch::Tensor& q,
+                          torch::Tensor& k,
+                          torch::Tensor& v,
+                          const int32_t rotary_dim,
+                          const float theta_base,
+                          torch::Tensor& batch_metadata,
+                          torch::Tensor& seq_metadata,
+                          torch::Tensor& tokens_to_seq,
+                          torch::Tensor& kv_ptrs);
+
+/*
+Copy into linear KV cache.
+*/
+void linear_kv_copy(torch::Tensor& kv_cache,
+                    torch::Tensor& q,
+                    torch::Tensor& k,
+                    torch::Tensor& v,
+                    torch::Tensor& batch_metadata,
+                    torch::Tensor& seq_metadata,
+                    torch::Tensor& tokens_to_seq,
+                    torch::Tensor& kv_ptrs);