chore: import upstream snapshot with attribution

deepspeed/inference/v2/modules/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from . import implementations
+from . import interfaces
+from .module_registry import ConfigBundle

deepspeed/inference/v2/modules/configs/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .attention_configs import (
+    DSSelfAttentionConfig,
+    PositionalEmbeddingType,
+    MaskingType,
+    RotateHalfConfig,
+)
+from .embedding_config import DSEmbeddingsConfig
+from .linear_config import DSLinearConfig
+from .moe_config import DSMoEConfig
+from .norm_config import DSNormConfig, NormTypeEnum
+from .unembed_config import DSUnembedConfig

deepspeed/inference/v2/modules/configs/attention_configs.py

@@ -0,0 +1,110 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from enum import Enum
+from typing import Dict, Optional
+
+from ...inference_utils import DtypeEnum
+from ...modules.ds_module import DSModuleConfig
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+
+
+class PositionalEmbeddingType(Enum):
+
+    # No positional embeddings
+    none = "none"
+
+    # Rotary positional embeddings - every half
+    rotate_half = "rotate_half"
+
+    # Rotary positional embeddings - every other
+    rotate_every_other = "rotate_every_other"
+
+    # Alibi
+    alibi = "alibi"
+
+
+class RotateHalfConfig(DeepSpeedConfigModel):
+
+    use_trained_freqs: bool = False
+    """
+    Whether to use a passed `trained_freqs` tensor for the attention implementation
+    or to use default synthesized frequencies.
+    """
+
+    theta_base: float = 10_000.0
+    """
+    Base for theta. This will only be used if `use_trained_freqs` is False.
+    """
+
+    rotate_dim: Optional[int] = None
+    """
+    How many neurons to rotate. If None, then all neurons will be rotated. Many external configs
+    will set this number to half the head dimension and then internally multiply by 2. To make it
+    more clear to understand what is happening (rotate_dim < head_dim -> then only partial rotation),
+    we do not do this multiplication internally.
+    """
+
+
+class MaskingType(Enum):
+
+    # No masking
+    none = "none"
+
+    # Causal masking
+    causal = "causal"
+
+    # Local masking
+    local = "local"
+
+    # Symmetric masking (this is a 1D tensor mask)
+    symmetric = "symmetric"
+
+    # Arbitrary masking (this would correspond to a 2D tensor mask)
+    asymmetric = "asymmetric"
+
+
+class DSSelfAttentionConfig(DSModuleConfig):
+    """
+    Config class for attention.
+    """
+
+    # Number of query attention heads on this shard
+    n_heads_q: int
+
+    # Number of KV attention heads on this shard
+    n_heads_kv: int
+
+    # Size of each attention head
+    head_size: int
+
+    # Max number of sequences that may compose a ragged batch
+    max_sequences: int
+
+    # Scale factor for attention scores
+    scale_factor: float = 1.0
+
+    # Input data type
+    input_dtype: DtypeEnum = DtypeEnum.fp16
+
+    # Output data type
+    output_dtype: DtypeEnum = DtypeEnum.fp16
+
+    # Masking type
+    masking_type: MaskingType = MaskingType.causal
+
+    # Masking args
+    masking_args: Dict = {}
+
+    # Positional embedding type
+    positional_embedding_type: PositionalEmbeddingType = PositionalEmbeddingType.none
+
+    # Positional embedding args
+    positional_embedding_config: Optional[RotateHalfConfig] = None
+    """
+    To extend this for the other positional embedding types, we would need to add
+    new configs for each type (as necessary) and annotate this with the
+    Union[RotateHalfConfig, OtherConfig, ...] type.
+    """

deepspeed/inference/v2/modules/configs/embedding_config.py

@@ -0,0 +1,70 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional
+
+from ...inference_utils import DtypeEnum, NormTypeEnum
+from ...modules.ds_module import DSModuleConfig
+"""
+Trying to define the space we need to support here right now:
+
+Types of embeddings I've found so far:
+    1. Token embedding
+    2. Position embedding
+    3. Token type embedding
+    4. LN
+
+GPTNeo: 1, 2, 3 (shared with 1)
+GPTNeoX: 1
+GPTJ: 1, 3
+LLaMA: 1
+BERT: 1, 2, 3, 4
+GPT2: 1, 2, 3 (shared with 1)
+
+Sidebar for OPT:
+OPT: 1, 2
+1 may not actually project to the actual hidden dimension according to the raw
+code, but for the model configs we care about it does.
+2 has a weird offset associated with it that the others do not.
+"""
+
+
+class DSEmbeddingsConfig(DSModuleConfig):
+    """
+    Config class for DSEmbeddings.
+    """
+
+    residual_dtype: DtypeEnum = DtypeEnum.fp16
+    """
+    Data type the module should use for its output.
+    """
+
+    embedding_dim: int
+    """
+    Dimensionality of the embedding projections.
+    """
+
+    positional_embedding: bool = False
+    """
+    Whether the module should expect a positional embedding matrix. The shape of this
+    matrix should be of shape [max_seq_len + positional_offset, embedding_dim]
+    """
+
+    positional_offset: int = 0
+    """
+    Whether the linearized token IDs should be offset by a certain amount. For an example
+    of this, see the OPT model implementation.
+    """
+
+    use_token_type: bool = False
+    """
+    Whether the module should expect a token type embedding matrix.
+    """
+
+    output_normalization: Optional[NormTypeEnum] = None
+    """
+    If a the output of the embedding module should be normalized, specify here. See
+    ``inference.inference_utils.NormTypeEnum`` for supported values.
+    """

deepspeed/inference/v2/modules/configs/linear_config.py

@@ -0,0 +1,43 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from ...inference_utils import ActivationType, DtypeEnum
+from ...modules.ds_module import DSModuleConfig
+
+
+class DSLinearConfig(DSModuleConfig):
+    """
+    Config class for DSLinearBase.
+    """
+
+    in_channels: int
+    """
+    Number of input channels
+    """
+
+    out_channels: int
+    """
+    Number of output channels. NOTE: If this linear layer is using a gated activation function,
+    the value for ``out_channels`` passed here should refer to the number of channels after
+    gating (i.e., the expected weight shape before transformations will be ``[out_channels * 2, in_channels]``).
+    """
+
+    activation: ActivationType = ActivationType.IDENTITY
+    """
+    The activation function for this layer. See :class:`deepspeed.inference.inference_utils.ActivationType` for
+    supported activation functions.
+    """
+
+    input_dtype: DtypeEnum = DtypeEnum.fp16
+    """
+    The data type of the input tensor. See :class:`deepspeed.inference.inference_utils.DtypeEnum` for supported
+    data types.
+    """
+
+    output_dtype: DtypeEnum = DtypeEnum.fp16
+    """
+    The data type of the output tensor. See :class:`deepspeed.inference.inference_utils.DtypeEnum` for supported
+    data types.
+    """

deepspeed/inference/v2/modules/configs/moe_config.py

@@ -0,0 +1,56 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from ...inference_utils import ActivationType, DtypeEnum
+from ...modules.ds_module import DSModuleConfig
+
+
+class DSMoEConfig(DSModuleConfig):
+    """
+    Config class for DSMoEBase
+    """
+
+    model_dim: int
+    """
+    Size of input activation.
+    """
+
+    intermediate_features: int
+    """
+    Size of intermediate activation. Specifically, this is the number of input features
+    in the second linear layer. Depending on the activation function, the output of the first
+    linear layer may have increased dimensionality.
+    """
+
+    n_experts: int
+    """
+    Number of experts.
+    """
+
+    top_k: int = 1
+    """
+    top-k gating function (like top-1 or top-2)
+    """
+
+    input_dtype: DtypeEnum = DtypeEnum.fp16
+    """
+    Data type for the input activations.
+    """
+
+    output_dtype: DtypeEnum = DtypeEnum.fp16
+    """
+    Data type for the output activations.
+    """
+
+    activation: ActivationType = ActivationType.IDENTITY
+    """
+    Activation function of the first MLP1
+    """
+
+    normalize_scores: bool = False
+    """
+    Whether normalization is applied to the selected scores. If true, the module
+    should rescale the scores such that their sum is 1.0.
+    """

deepspeed/inference/v2/modules/configs/norm_config.py

@@ -0,0 +1,32 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from ...inference_utils import DtypeEnum, NormTypeEnum
+from ...modules.ds_module import DSModuleConfig
+
+
+class DSNormConfig(DSModuleConfig):
+    """
+    Config class for both DSPreLN and DSPostLN.
+    """
+
+    # Type of normalization
+    type: NormTypeEnum
+
+    # Number of channels in the model embedding
+    channels: int
+
+    # Data type of the residual input/outputs (we assume the residual must
+    # be the same data type for the entire model).
+    residual_dtype: DtypeEnum = DtypeEnum.fp16
+
+    # Data type of the hidden states input
+    input_dtype: DtypeEnum = DtypeEnum.fp16
+
+    # Data type of the hidden states output
+    output_dtype: DtypeEnum = DtypeEnum.fp16
+
+    # Epsilon value for numerical stability
+    eps: float = 1e-5

deepspeed/inference/v2/modules/configs/unembed_config.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from ...inference_utils import DtypeEnum, NormTypeEnum
+from ...modules.ds_module import DSModuleConfig
+from typing import Optional
+
+
+class DSUnembedConfig(DSModuleConfig):
+    """
+    Config class for DSUnembed
+    """
+
+    dtype: DtypeEnum = DtypeEnum.fp16
+    """
+    Expected data type.
+    """
+
+    norm_type: Optional[NormTypeEnum] = None
+    """
+    Whether the input to the unembed is normalized prior to the unembedding projection.
+    """
+
+    model_dim: int
+    """
+    Model embedding size.
+    """
+
+    max_sequences: int
+    """
+    Max sequences composing the ragged batch.
+    """
+
+    vocab_size: int
+    """
+    Local vocab size (the full vocab size may have been sharded across model parallel ranks)
+    """

deepspeed/inference/v2/modules/ds_module.py

@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import ABC, abstractstaticmethod
+from typing import Any, Dict, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+
+
+class DSModuleConfig(DeepSpeedConfigModel):
+
+    max_tokens: int
+
+
+class DSModuleBase(torch.nn.Module, ABC):
+    """
+    Base class for all DeepSpeed Inference modules. This class establishes
+    the basic attributes of a DSModule. Only abstract functionality modules should inherit
+    directly from this class, not specific implementations.
+    """
+
+    @abstractstaticmethod
+    def name() -> str:
+        """
+        Return a memorable, human-readable name for this module.
+
+        This will be used as a key in custom inference configurations and should only
+        be implemented by the children of functionality modules.
+        """
+        ...
+
+    @abstractstaticmethod
+    def config_class() -> Type[DSModuleConfig]:
+        """
+        Return the associated config class for this module.
+
+        This should be implemented (along with the config class) by an abstract functionality
+        module.
+        """
+        ...
+
+    @abstractstaticmethod
+    def supports_config(config: DSModuleConfig) -> bool:
+        """
+        Return whether or not this module supports the given config.
+
+        This should be implemented by the children of functionality modules and should report
+        whether it would be feasible to instantiate this module with the given config.
+        """
+        ...
+
+    def __init__(self, config: DSModuleConfig, implementation_config: Dict[str, Any] = {}) -> None:
+        """
+        Initialize the module with the given config.
+        """
+        super().__init__()
+        self._config = config
+        self._implementation_config = implementation_config

deepspeed/inference/v2/modules/heuristics.py

@@ -0,0 +1,195 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from ..config_v2 import RaggedInferenceEngineConfig
+from ..inference_utils import NormTypeEnum
+
+from .module_registry import ConfigBundle
+from ..modules.configs import (
+    DSEmbeddingsConfig,
+    DSLinearConfig,
+    DSMoEConfig,
+    DSNormConfig,
+    DSSelfAttentionConfig,
+    DSUnembedConfig,
+)
+from ..modules.interfaces import (
+    DSEmbeddingBase,
+    DSEmbeddingRegistry,
+    DSLinearBase,
+    DSLinearRegistry,
+    DSMoEBase,
+    DSMoERegistry,
+    DSPostNormBase,
+    DSPostNormRegistry,
+    DSPreNormBase,
+    DSPreNormRegistry,
+    DSSelfAttentionBase,
+    DSSelfAttentionRegistry,
+    DSUnembedBase,
+    DSUnembedRegistry,
+)
+
+
+def instantiate_attention(attention_config: DSSelfAttentionConfig,
+                          engine_config: RaggedInferenceEngineConfig) -> DSSelfAttentionBase:
+    """
+    Choose an appropriate attention implementation based on the given configurations. This
+    method is currently a stub, but as more implementations may be developed  we can centralize
+    the logic for choosing between them here.
+
+    Arguments:
+        attention_config (DSSelfAttentionConfig): Configuration for the attention module.
+        engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine.
+
+    Returns:
+        An attention module implementing the given configuration.
+    """
+
+    # Currently, we only have one implementation, so we just return it.
+    config = ConfigBundle(name="dense_blocked_attention", config=attention_config)
+    return DSSelfAttentionRegistry.instantiate_config(config)
+
+
+def instantiate_embed(embed_config: DSEmbeddingsConfig, engine_config: RaggedInferenceEngineConfig) -> DSEmbeddingBase:
+    """
+    Choose an appropriate embedding implementation based on the given configurations. This
+    method is currently a stub, but as more implementations may be developed  we can centralize
+    the logic for choosing between them here.
+
+    Arguments:
+        embed_config (DSEmbeddingsConfig): Configuration for the embedding module.
+        engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine.
+
+    Returns:
+        An embedding module implementing the given configuration.
+    """
+
+    # Currently, we only have one implementation, so we just return it.
+    config = ConfigBundle(name="ragged_embedding", config=embed_config)
+    return DSEmbeddingRegistry.instantiate_config(config)
+
+
+def instantiate_linear(linear_config: DSLinearConfig, engine_config: RaggedInferenceEngineConfig) -> DSLinearBase:
+    """
+    Choose an appropriate linear implementation based on the given configurations. This
+    method is currently a stub, but as more implementations may be developed  we can centralize
+    the logic for choosing between them here.
+
+    Arguments:
+        linear_config (DSLinearConfig): Configuration for the linear module.
+        engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine.
+
+    Returns:
+        A linear module implementing the given configuration.
+    """
+
+    quantization_mode = engine_config.quantization.quantization_mode
+    if quantization_mode is None:
+        config = ConfigBundle(name="blas_fp_linear", config=linear_config)
+    else:
+        # Currently, we only support ``quantized_wf6af16_linear`` on NVIDIA Ampere GPUs.
+        if quantization_mode == "wf6af16":
+            import torch
+            if not torch.cuda.is_available():  #ignore-cuda
+                raise ValueError("WF6AF16 quantization is only supported on CUDA")
+            else:
+                is_rocm_pytorch = hasattr(torch.version, 'hip') and torch.version.hip is not None
+                if is_rocm_pytorch:
+                    raise ValueError("WF6AF16 quantization is only supported on NVIDIA GPUs")
+                elif torch.cuda.get_device_properties(0).major != 8:  #ignore-cuda
+                    raise ValueError("WF6AF16 quantization is only supported on Ampere architectures")
+            config = ConfigBundle(name="quantized_wf6af16_linear", config=linear_config)
+        else:
+            raise ValueError(f"Unsupported quantization mode: {quantization_mode}")
+    return DSLinearRegistry.instantiate_config(config)
+
+
+def instantiate_moe(moe_config: DSMoEConfig, engine_config: RaggedInferenceEngineConfig) -> DSMoEBase:
+    """
+    Choose an appropriate MoE implementation based on the given configurations. This
+    method is currently a stub, but as more implementations may be developed  we can centralize
+    the logic for choosing between them here.
+
+    Arguments:
+        moe_config (DSMoEConfig): Configuration for the MoE module.
+        engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine.
+
+    Returns:
+        A MoE module implementing the given configuration.
+    """
+
+    moe_type = "cutlass_multi_gemm_moe"
+
+    if moe_type == "cutlass_multi_gemm_moe":
+        # TODO: Get this off an engine config
+        implementation_config = {
+            "weight_dtype": moe_config.input_dtype,
+        }
+
+    # Currently, we only have one implementation, so we just return it.
+    config = ConfigBundle(name="cutlass_multi_gemm_moe",
+                          config=moe_config,
+                          implementation_config=implementation_config)
+    return DSMoERegistry.instantiate_config(config)
+
+
+def instantiate_post_norm(norm_config: DSNormConfig, engine_config: RaggedInferenceEngineConfig) -> DSPostNormBase:
+    """
+    Choose an appropriate post-norm implementation based on the given configurations. This
+    method is currently a stub, but as more implementations may be developed  we can centralize
+    the logic for choosing between them here.
+
+    Arguments:
+        norm_config (DSNormConfig): Configuration for the post-norm module.
+        engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine.
+
+    Returns:
+        A post-norm module implementing the given configuration.
+    """
+
+    # Currently, we only have one implementation, so we just return it.
+    config = ConfigBundle(name="cuda_post_ln", config=norm_config)
+    return DSPostNormRegistry.instantiate_config(config)
+
+
+def instantiate_pre_norm(norm_config: DSNormConfig, engine_config: RaggedInferenceEngineConfig) -> DSPreNormBase:
+    """
+    Choose an appropriate pre-norm implementation based on the given configurations. Currently,
+    this will select between two CUDA implementations, one for LayerNorm and one for RMSNorm.
+
+    Arguments:
+        norm_config (DSNormConfig): Configuration for the pre-norm module.
+        engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine.
+
+    Returns:
+        A pre-norm module implementing the given configuration.
+    """
+    if NormTypeEnum(norm_config.type) == NormTypeEnum.LayerNorm:
+        module_name = "cuda_pre_ln"
+    elif NormTypeEnum(norm_config.type) == NormTypeEnum.RMSNorm:
+        module_name = "cuda_pre_rms"
+
+    config = ConfigBundle(name=module_name, config=norm_config)
+    return DSPreNormRegistry.instantiate_config(config)
+
+
+def instantiate_unembed(unembed_config: DSUnembedConfig, engine_config: RaggedInferenceEngineConfig) -> DSUnembedBase:
+    """
+    Choose an appropriate unembedding implementation based on the given configurations. This
+    method is currently a stub, but as more implementations may be developed  we can centralize
+    the logic for choosing between them here.
+
+    Arguments:
+        unembed_config (DSUnembedConfig): Configuration for the unembed module.
+        engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine.
+
+    Returns:
+        An unembed module implementing the given configuration.
+    """
+
+    # Currently, we only have one implementation, so we just return it.
+    config = ConfigBundle(name="ragged_unembed", config=unembed_config)
+    return DSUnembedRegistry.instantiate_config(config)

deepspeed/inference/v2/modules/implementations/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# Imports for registering ops
+from .attention import *
+from .linear import *
+from .post_norm import *
+from .pre_norm import *
+from .embedding import *
+from .unembed import *
+from .moe import *

deepspeed/inference/v2/modules/implementations/attention/__init__.py

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

deepspeed/inference/v2/modules/implementations/attention/dense_blocked_attention.py

@@ -0,0 +1,180 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from ....allocator import empty_from
+from ....inference_utils import DtypeEnum
+from ....kernels.ragged_ops import (
+    AtomBuilder,
+    BlockedFlashAttn,
+    BlockedRotaryEmbeddings,
+    BlockedTrainedRotaryEmbeddings,
+    get_q_block_size,
+    get_kv_block_size,
+    LinearBlockedKVCopy,
+)
+from ....ragged import RaggedBatchWrapper, split_kv
+from deepspeed.ops.op_builder import RaggedUtilsBuilder
+
+from ...interfaces import DSSelfAttentionBase, DSSelfAttentionRegistry
+from ...configs import DSSelfAttentionConfig, PositionalEmbeddingType, MaskingType
+
+try:
+    from functools import cached_property
+except ImportError:
+
+    def cached_property(func):
+        return property(func)
+
+
+@DSSelfAttentionRegistry.register_module
+class DSDenseBlockedAttention(DSSelfAttentionBase):
+    """
+    Self attention implementation for dense, blocked self attention.
+    """
+
+    @staticmethod
+    def name() -> str:
+        return 'dense_blocked_attention'
+
+    @staticmethod
+    def supports_config(config: DSSelfAttentionConfig) -> bool:
+
+        if config.input_dtype != config.output_dtype:
+            return False
+
+        if DtypeEnum(config.input_dtype) not in (DtypeEnum.fp16, DtypeEnum.bf16):
+            return False
+
+        if PositionalEmbeddingType(config.positional_embedding_type) not in [
+                PositionalEmbeddingType.none, PositionalEmbeddingType.rotate_half
+        ]:
+            return False
+
+        if MaskingType(config.masking_type) != MaskingType.causal:
+            return False
+
+        return True
+
+    def __init__(self, config: DSSelfAttentionConfig, implementation_config: Dict[str, Any]) -> None:
+        """
+        Create the Attention DSModule.
+
+        Args:
+            config (DSSelfAttentionConfig): The self attention config for all attention DSModules.
+            implementation_config (Dict[str, Any]):
+                There are two (dependent) potential components in the implementtion config.
+
+                1. `trained_freqs` - If the embedding weights for RoPE are trained, the implementation
+                config should contain {'trained_freqs': True}. This will mean the implementation will
+                expect a `trained_freqs` tensor in the `forward` method and will not synthesize the
+                values internally.
+
+                2. `theta_base` - The base value for synthesized frequencies in the rotary embeddings.
+                This will only be used if `trained_freqs` is False or not present in the `implementation_config`. If this is not included, the default value of 10000.0 will be used.
+        """
+        super().__init__(config, implementation_config)
+
+        embed_type = PositionalEmbeddingType(config.positional_embedding_type)
+        if embed_type == PositionalEmbeddingType.none:
+            self._kv_copy = LinearBlockedKVCopy(self._config.head_size, self._config.n_heads_q,
+                                                self._config.n_heads_kv, self._config.input_dtype)
+        elif embed_type == PositionalEmbeddingType.rotate_half:
+            rotary_config = config.positional_embedding_config
+            assert rotary_config is not None, "Rotary config must be provided if using rotate_half as Positional Embedding Type."
+
+            if rotary_config.use_trained_freqs:
+                # Theta and rotary dim are effectively embedded into either the values (theta) or the shape (rotary_dim)
+                # of the trained_freqs tensor.
+                self._kv_copy = BlockedTrainedRotaryEmbeddings(self._config.head_size, self._config.n_heads_q,
+                                                               self._config.n_heads_kv, self._config.input_dtype)
+            else:
+                theta_base = rotary_config.theta_base
+                rotary_dim = rotary_config.rotate_dim if rotary_config.rotate_dim is not None else self._config.head_size
+                self._kv_copy = BlockedRotaryEmbeddings(self._config.head_size, self._config.n_heads_q,
+                                                        self._config.n_heads_kv, self._config.input_dtype, rotary_dim,
+                                                        theta_base)
+
+        self._softmax_scale = self._config.scale_factor
+
+        # TODO(cmikeh2): Attention kernel gets created here.
+        self._attn_kernel = BlockedFlashAttn(self._config.head_size, self._config.input_dtype)
+        self._atom_builder = AtomBuilder()
+
+        self.model_dim = self._config.head_size * self._config.n_heads_q
+        self._output = torch.empty((self._config.max_tokens, self._config.head_size * self._config.n_heads_q),
+                                   dtype=self._config.output_dtype,
+                                   device=get_accelerator().current_device())
+
+        # TODO(cmikeh2): Pre-allocate storage buffer for the attention atoms.
+        self._max_atoms = self._config.max_sequences
+        self._atoms = torch.empty((self._max_atoms, 8), dtype=torch.int32, device=get_accelerator().current_device())
+
+        alloc_func = RaggedUtilsBuilder().load().allocate_fast_host_buffer
+        self._atoms_shadow = alloc_func(self._atoms)
+        self._cur_atoms = 0
+
+    @cached_property
+    def kv_block_size(self) -> int:
+        """
+        Return preferred granulatity for blocked KV-cache implementation.
+        """
+        return get_kv_block_size(self._config.head_size)
+
+    @cached_property
+    def q_block_size(self) -> int:
+        """
+        Property to calculate blocking granularity for the query dimension.
+        This has no impact on the KV-cache structure, but will  affect the
+        number of attention atoms associated with a batch.
+        """
+        return get_q_block_size(self._config.head_size)
+
+    def build_atoms(self, ragged_batch: RaggedBatchWrapper) -> None:
+        """
+        Build the atoms for the attention kernel.
+
+        Args:
+            ragged_batch (RaggedBatchWrapper): The input ids and associated ragged batch metadata.
+        """
+        host_atoms, n_atoms = self._atom_builder(self._atoms_shadow, ragged_batch, self.q_block_size,
+                                                 self.kv_block_size)
+
+        self._cur_atoms = n_atoms
+        self._atoms[:n_atoms].copy_(host_atoms[:n_atoms], non_blocking=True)
+
+    def forward(self,
+                q_k_v: torch.Tensor,
+                kv_cache: torch.Tensor,
+                batch: RaggedBatchWrapper,
+                inv_freqs: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Forward implementation.
+
+        Args:
+            q_k_v (torch.Tensor): Query/Key/Value projection Tensor of shape
+                [n_heads, (n_heads_q + 2 * n_heads_kv) * head_size].
+            kv_cache (torch.Tensor): Blocked persistent cache of shape
+                [2, batch, block_size, n_heads_kv, head_size].
+            batch (RaggedBatchWrapper): The input ids and associated ragged batch metadata.
+            inv_freqs (Optional[torch.Tensor]): The inverse frequencies for the rotary embeddings if they
+                have been modified from synthesizable values.
+        """
+        if inv_freqs is not None:
+            self._kv_copy(kv_cache, q_k_v, batch, inv_freqs)
+        else:
+            self._kv_copy(kv_cache, q_k_v, batch)
+
+        q = q_k_v[:, :self._config.head_size * self._config.n_heads_q]
+        output = empty_from(self._output, q.shape)
+        k_cache, v_cache = split_kv(kv_cache)
+
+        self._attn_kernel(output, q, k_cache, v_cache, self._atoms[:self._cur_atoms], self._softmax_scale)
+
+        return output

deepspeed/inference/v2/modules/implementations/embedding/__init__.py

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

deepspeed/inference/v2/modules/implementations/embedding/ragged_embedding.py

@@ -0,0 +1,77 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from ....allocator import empty_from
+from ....inference_utils import DtypeEnum
+from ....kernels.ragged_ops import RaggedEmbeddingKernel
+from ....ragged import RaggedBatchWrapper
+from ...interfaces import DSEmbeddingBase, DSEmbeddingRegistry
+from ...configs import DSEmbeddingsConfig
+
+
+@DSEmbeddingRegistry.register_module
+class DSRaggedEmbedding(DSEmbeddingBase):
+
+    @staticmethod
+    def name():
+        return 'ragged_embedding'
+
+    @staticmethod
+    def supports_config(config: DSEmbeddingsConfig) -> bool:
+
+        if DtypeEnum(config.residual_dtype) not in [DtypeEnum.fp16, DtypeEnum.bf16, DtypeEnum.fp32]:
+            return False
+
+        if config.use_token_type:
+            return False
+
+        if config.output_normalization is not None:
+            return False
+
+        try:
+            _ = RaggedEmbeddingKernel(config.residual_dtype, torch.int32, config.embedding_dim)
+        except ValueError:
+            return False
+
+        return True
+
+    def __init__(self, config: DSEmbeddingsConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+        self.embed_offset = self._config.positional_offset
+
+        # TODO(cmikeh2): How do we want to avoid the int32 vs int64 issue?
+        self._ragged_embed = RaggedEmbeddingKernel(self._config.residual_dtype, torch.int32,
+                                                   self._config.embedding_dim)
+
+        self._output = torch.empty((self._config.max_tokens, self._config.embedding_dim),
+                                   dtype=self._config.residual_dtype,
+                                   device=get_accelerator().current_device())
+
+    @property
+    def output(self) -> torch.Tensor:
+        return self._output
+
+    def forward(self,
+                ragged_batch: RaggedBatchWrapper,
+                word_embeddings: torch.Tensor,
+                position_embeddings: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Parameters:
+            ragged_batch (RaggedBatchWrapper): The input ids and associated ragged batch metadata.
+            word_embeddings (torch.Tensor): The word embedding table
+        """
+        output = empty_from(self._output, (ragged_batch.tensor_toks, self._config.embedding_dim))
+        self._ragged_embed(output,
+                           ragged_batch,
+                           word_embeddings,
+                           position_embed_weight=position_embeddings,
+                           position_embed_offset=self.embed_offset)
+        return output

deepspeed/inference/v2/modules/implementations/linear/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .blas_fp_linear import BlasFPLinear
+from .quantized_linear import QuantizedWf6Af16Linear, fp_quantize

deepspeed/inference/v2/modules/implementations/linear/blas_fp_linear.py

@@ -0,0 +1,103 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from ....allocator import empty_from
+from ....inference_utils import is_gated
+from ....kernels.core_ops import (
+    BlasLibLinear,
+    CUDABiasActivation,
+    CUDAGatedActivation,
+)
+
+from ...interfaces import DSLinearBase, DSLinearRegistry
+from ...configs import DSLinearConfig
+from ....inference_parameter import InferenceParameter
+
+
+@DSLinearRegistry.register_module
+class BlasFPLinear(DSLinearBase):
+    """
+    Linear DSModule based on BLAS library and standalone bias + activation kernel implementation.
+    """
+
+    @staticmethod
+    def name():
+        return 'blas_fp_linear'
+
+    @staticmethod
+    def supports_config(config: DSLinearConfig) -> bool:
+        if config.input_dtype != config.output_dtype:
+            return False
+
+        if config.input_dtype != torch.float16 and config.input_dtype != torch.bfloat16:
+            return False
+
+        if is_gated(config.activation):
+            try:
+                _ = CUDAGatedActivation(config.out_channels, config.output_dtype, config.activation)
+            except ValueError:
+                return False
+        else:
+            try:
+                _ = CUDABiasActivation(config.out_channels, config.output_dtype, config.activation)
+            except ValueError:
+                return False
+
+        return True
+
+    def __init__(self, config: DSLinearConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+        self._linear_impl = BlasLibLinear(self._config.input_dtype)
+
+        if is_gated(config.activation):
+            self._is_gated = True
+            self._act_fn = CUDAGatedActivation(config.out_channels, config.output_dtype, config.activation)
+            self._double_buffer = torch.empty((config.max_tokens, config.out_channels * 2),
+                                              dtype=config.output_dtype,
+                                              device=get_accelerator().current_device())
+        else:
+            self._is_gated = False
+            self._act_fn = CUDABiasActivation(config.out_channels, config.output_dtype, config.activation)
+
+        self._output = torch.empty((config.max_tokens, config.out_channels),
+                                   dtype=config.output_dtype,
+                                   device=get_accelerator().current_device())
+
+    def transform_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Converts param to same data type as input and output.
+
+        Parameters:
+            param (torch.Tensor): Weight or bias tensor.
+        """
+        param = param.to(self._config.output_dtype)
+        return InferenceParameter.initialize(param)
+
+    def forward(self, hidden_states: torch.Tensor, w: torch.Tensor, b: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        output = empty_from(self._output, (hidden_states.shape[0], self._config.out_channels))
+
+        if self._is_gated:
+            staging_output = empty_from(self._double_buffer, (hidden_states.shape[0], self._config.out_channels * 2))
+            self._linear_impl(staging_output, hidden_states, w)
+            self._act_fn(output, staging_output, b)
+        else:
+            self._linear_impl(output, hidden_states, w)
+            self._act_fn(output, b)
+
+        return output
+
+    @property
+    def output(self) -> torch.Tensor:
+        """
+        Return the padded, pre-allocated output Tensor.
+        """
+        return self._output

deepspeed/inference/v2/modules/implementations/linear/quantized_linear.py

@@ -0,0 +1,205 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.ops.op_builder import InferenceCoreBuilder
+from ....allocator import empty_from
+from ....inference_utils import is_gated
+from ....kernels.core_ops import (
+    CUDAWf6Af16Linear,
+    CUDABiasActivation,
+    CUDAGatedActivation,
+)
+
+from ...interfaces import DSLinearBase, DSLinearRegistry
+from ...configs import DSLinearConfig
+from ....inference_parameter import InferenceParameter
+
+
+def fp_quantize(input: torch.FloatTensor,
+                num_bits: int = 6,
+                exp_bits: int = 3,
+                min_value: torch.FloatTensor = None,
+                max_value: torch.FloatTensor = None,
+                group_size: int = -1):
+    """
+    Args:
+        inputs (`torch.FloatTensor`)
+            The input which needs to be quantized
+        num_bits (int, >=4)
+            Number of bits to use for quantization
+        exp_bits:
+            fp exp_bits
+        min_value/max_vlue (torch.FloatTensor)
+            Used for static activation quantization
+        group_size (int) N
+            The quantization block size, each N numbers has its own scaling
+            factor and off-site. -1 means use the last dim as the group_size
+    Returns:
+        quantized_fake_fp6
+            The quantized weights, in fp16 format and contains fp6 value.
+        scales
+            Quantization scales
+    """
+
+    try:
+        from qtorch.quant import float_quantize
+    except ImportError:
+        raise ImportError("Please install qtorch to use this function")
+
+    assert (min_value is None and max_value is None) or (min_value is not None and max_value is not None)
+
+    assert input.dtype == torch.float16
+
+    orig_device = input.device
+    input = input.to(torch.float32).to(get_accelerator().current_device())
+    if num_bits == 6 and exp_bits == 3:  # this is default
+        q_range = 28
+    else:
+        raise NotImplementedError
+
+    man_bits = num_bits - exp_bits - 1
+    input_shape = input.shape
+
+    if group_size == -1:
+        group_size = input_shape[-1]
+    else:
+        # Only support per-channel quantization
+        raise NotImplementedError
+    num_groups = input.numel() // group_size
+    input = input.reshape(num_groups, -1)
+
+    if min_value is None:
+        max_input = torch.amax(torch.abs(input), dim=-1).view(num_groups, -1)
+    else:
+        max_input = torch.max(min_value.abs(), max_value)  # .view(-1)
+    scales = max_input / q_range  # q_range + 1
+    scales[scales == 0] = 1  # avoid zero scales
+    scaled_input = input / scales
+
+    quantized_fake_fp6 = float_quantize(scaled_input, exp_bits, man_bits, rounding="nearest")
+
+    quantized_fake_fp6 = quantized_fake_fp6.reshape(input_shape).contiguous().to(torch.float16).to(orig_device)
+    scales = scales.to(torch.float16).to(orig_device)
+    # Now the dequantized value is quantized_fake_fp6 * scales
+
+    return quantized_fake_fp6, scales
+
+
+@DSLinearRegistry.register_module
+class QuantizedWf6Af16Linear(DSLinearBase):
+    """
+    Linear DSModule for FP6 weight-only quantization kernel, where weight is FP6
+    and activation is FP16.
+    """
+
+    @staticmethod
+    def name():
+        return 'quantized_wf6af16_linear'
+
+    @staticmethod
+    def supports_config(config: DSLinearConfig) -> bool:
+        if config.input_dtype != config.output_dtype:
+            return False
+
+        # As for fp6 data items, they are packed and stored in a set of fp16
+        # tensors. E.g., 8 fp6 data items are stored in 3 fp16 tensor.
+        if config.input_dtype != torch.float16:
+            return False
+
+        if is_gated(config.activation):
+            try:
+                _ = CUDAGatedActivation(config.out_channels, config.output_dtype, config.activation)
+            except ValueError:
+                return False
+        else:
+            try:
+                _ = CUDABiasActivation(config.out_channels, config.output_dtype, config.activation)
+            except ValueError:
+                return False
+
+        return True
+
+    def __init__(self, config: DSLinearConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+        self._linear_impl = CUDAWf6Af16Linear()
+
+        if is_gated(config.activation):
+            # In the FP6 kernel implementation, the MatMul is W * A, where W is
+            # the weight and A is activation. M is the output channel size.
+            self.out_channels = self._config.out_channels * 2
+            self.in_channels = self._config.in_channels
+            self._is_gated = True
+            self._act_fn = CUDAGatedActivation(config.out_channels, config.output_dtype, config.activation)
+            self._double_buffer = torch.empty((config.max_tokens, config.out_channels * 2),
+                                              dtype=config.output_dtype,
+                                              device=get_accelerator().current_device())
+        else:
+            self.out_channels = self._config.out_channels
+            self.in_channels = self._config.in_channels
+            self._is_gated = False
+            self._act_fn = CUDABiasActivation(config.out_channels, config.output_dtype, config.activation)
+
+        self._output = torch.empty((config.max_tokens, config.out_channels),
+                                   dtype=config.output_dtype,
+                                   device=get_accelerator().current_device())
+
+        self.inf_module = InferenceCoreBuilder().load()
+        self.inf_module.create_handle()
+        self.preprocess_weight = self.inf_module.preprocess_weight
+
+        self.quantizer = fp_quantize
+
+    def transform_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Converts param to same data type as input and output.
+
+        Parameters:
+            param (torch.Tensor): Weight or bias tensor.
+        """
+        # It expects that the quantization scales are store in the attribute `scales`.
+
+        if param.ndim == 1:  # bias, do nothing
+            return InferenceParameter.initialize(param)
+
+        quantized_fake_fp6, scales = self.quantizer(param, num_bits=6, exp_bits=3)
+
+        # This is for debugging, will delete before release.
+        assert (quantized_fake_fp6.dtype == torch.float16)
+        assert quantized_fake_fp6.shape[0] == self.out_channels
+        assert scales.numel() == self.out_channels
+
+        weights_2bit, weights_4bit = self.preprocess_weight(quantized_fake_fp6)
+
+        return InferenceParameter.initialize(weights_2bit, weights_4bit=weights_4bit, scales=scales)
+
+    def forward(self, hidden_states: torch.Tensor, w: torch.Tensor, b: Optional[torch.Tensor] = None) -> torch.Tensor:
+        weights_2bit = w
+        weights_4bit = w.weights_4bit
+        scales = w.scales
+        output = empty_from(self._output, (hidden_states.shape[0], self._config.out_channels))
+        if self._is_gated:
+            staging_output = empty_from(self._double_buffer, (hidden_states.shape[0], self.out_channels))
+            self._linear_impl(staging_output, hidden_states, weights_2bit, weights_4bit, scales, self.out_channels,
+                              hidden_states.shape[0], self.in_channels)
+            self._act_fn(output, staging_output, b)
+        else:
+            self._linear_impl(output, hidden_states, weights_2bit, weights_4bit, scales, self.out_channels,
+                              hidden_states.shape[0], self.in_channels)
+            self._act_fn(output, b)
+
+        return output
+
+    @property
+    def output(self) -> torch.Tensor:
+        """
+        Return the padded, pre-allocated output Tensor.
+        """
+        return self._output

deepspeed/inference/v2/modules/implementations/moe/__init__.py

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

deepspeed/inference/v2/modules/implementations/moe/cutlass_multi_gemm.py

@@ -0,0 +1,249 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from ....allocator import empty_from
+from ....inference_utils import ActivationType, is_gated
+from ....kernels.core_ops import BlasLibLinear, CUDAGatedActivation
+from ....kernels.ragged_ops import (
+    MoEGather,
+    MoEScatter,
+    RaggedTopKGating,
+)
+from ....ragged import RaggedBatchWrapper
+
+from ...interfaces import DSMoEBase, DSMoERegistry
+from ...configs import DSMoEConfig
+from ....kernels.cutlass_ops import MoEGEMM
+from ....inference_parameter import InferenceParameter
+
+
+@DSMoERegistry.register_module
+class DSMultiGemmMoE(DSMoEBase):
+    """
+    MoE implementation based on the CUTLASS multi-GEMM.
+    """
+
+    @staticmethod
+    def name():
+        return 'cutlass_multi_gemm_moe'
+
+    @staticmethod
+    def supports_config(config: DSMoEConfig) -> bool:
+        if config.input_dtype != config.output_dtype:
+            return False
+
+        if config.input_dtype != torch.float16 and config.input_dtype != torch.bfloat16:
+            return False
+
+        if config.top_k != 1 and config.top_k != 2 and config.top_k != 4 and config.top_k != 8:
+            return False
+
+        return True
+
+    def __init__(self, config: DSMoEConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+        # Convenience variables for frequently accessed items.
+        self.max_tokens = self._config.max_tokens
+        self.n_experts = self._config.n_experts
+        self.n_top_k = self._config.top_k
+        self.intermediate_dim = self._config.intermediate_features
+
+        moe_op_act_fn = ActivationType.IDENTITY if is_gated(self._config.activation) else self._config.activation
+
+        self._mlp_1 = MoEGEMM(fp_dtype=implementation_config['weight_dtype'], act_fn=moe_op_act_fn)
+        self._mlp_2 = MoEGEMM(fp_dtype=implementation_config['weight_dtype'], act_fn=ActivationType.IDENTITY)
+
+        if is_gated(self._config.activation):
+            self._activation = CUDAGatedActivation(self._config.model_dim, self._config.input_dtype,
+                                                   self._config.activation)
+        else:
+            self._activation = None
+
+        self._gate_proj = BlasLibLinear(self._config.input_dtype)
+        self._top_1_gate = RaggedTopKGating(config.input_dtype)
+        self._moe_scatter = MoEScatter(config.input_dtype, config.model_dim)
+        self._moe_gather = MoEGather(config.input_dtype, config.model_dim, config.normalize_scores)
+
+        self._create_buffers()
+
+    def _create_buffers(self):
+
+        # Gating buffers
+        self._logits = torch.empty((self._config.max_tokens, self.n_experts),
+                                   dtype=self._config.input_dtype,
+                                   device=get_accelerator().current_device())
+        self._expert_counts = torch.empty((self.n_experts, ),
+                                          dtype=torch.int32,
+                                          device=get_accelerator().current_device())
+        self._scores = torch.empty((self._config.max_tokens, self.n_top_k),
+                                   dtype=torch.float32,
+                                   device=get_accelerator().current_device())
+        self._assignments = torch.empty((self._config.max_tokens, self.n_top_k),
+                                        dtype=torch.int32,
+                                        device=get_accelerator().current_device())
+        self._offsets = torch.empty((self._config.max_tokens, self.n_top_k),
+                                    dtype=torch.int32,
+                                    device=get_accelerator().current_device())
+
+        # Scatter buffers
+        self._moe_input = torch.empty((self._config.max_tokens * self.n_top_k, self._config.model_dim),
+                                      dtype=self._config.input_dtype,
+                                      device=get_accelerator().current_device())
+        self._expert_cumsum = torch.empty((self._config.n_experts, ),
+                                          dtype=torch.int64,
+                                          device=get_accelerator().current_device())
+        self._mapped_slots = torch.empty((self._config.max_tokens, self.n_top_k),
+                                         dtype=torch.int32,
+                                         device=get_accelerator().current_device())
+
+        # GEMM Buffers
+        self._intermediate = torch.empty((self._config.max_tokens * self.n_top_k, self._config.intermediate_features),
+                                         dtype=self._config.output_dtype,
+                                         device=get_accelerator().current_device())
+        if self._activation is not None:
+            self._gated_intermediate = torch.empty(
+                (self._config.max_tokens * self.n_top_k, self._config.intermediate_features * 2),
+                dtype=self._config.output_dtype,
+                device=get_accelerator().current_device())
+
+        self._output_unordered = torch.empty((self._config.max_tokens * self.n_top_k, self._config.model_dim),
+                                             dtype=self._config.output_dtype,
+                                             device=get_accelerator().current_device())
+
+        # Gather buffer
+        self._output = torch.empty((self._config.max_tokens, self._config.model_dim),
+                                   dtype=self._config.output_dtype,
+                                   device=get_accelerator().current_device())
+
+    def transform_gate_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Ensures gate param is going to match the activation data type.
+        """
+        param = param.to(self._config.input_dtype)
+        return InferenceParameter.initialize(param)
+
+    def transform_moe_mlp_1_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Converts param to same data type as input and output.
+
+        Parameters:
+            param (torch.Tensor): Weight or bias tensor.
+        """
+        param = param.to(self._config.input_dtype)
+
+        if len(param.shape) == 3:
+            param = param.permute(0, 2, 1).contiguous()
+        return InferenceParameter.initialize(param)
+
+    def transform_moe_mlp_2_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Converts param to same data type as input and output.
+
+        Parameters:
+            param (torch.Tensor): Weight or bias tensor.
+        """
+        param = param.to(self._config.input_dtype)
+
+        if len(param.shape) == 3:
+            param = param.permute(0, 2, 1).contiguous()
+        return InferenceParameter.initialize(param)
+
+    @property
+    def output(self) -> torch.Tensor:
+        return self._output
+
+    def _gate(self, hidden_states: torch.Tensor, batch_metadata: RaggedBatchWrapper,
+              gate_w: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Helper function to isolate the logit for gating. This will take the hidden states
+        and produce the metadata + tensors for the CUTLASS ragged GEMMs. If the input has
+        been padded for CG, this will strip the padding for MoE.
+
+        Parameters:
+            hidden_states (torch.Tensor): Hidden states tensor. Expected shape is [n_tokens, model_dim].
+            batch_metadata (RaggedBatchWrapper): Batch metadata for the hidden states.
+            gate_w (torch.Tensor): Gate weight tensor. Expected shape is [num_experts, model_dim].
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: The MoE input, the cumsum of the offsets (for the MoE kernels themselves), the scores, and the mapped slots (to recover the original order of the tokens)
+        """
+
+        # Get views on the buffers for gating
+        logits = empty_from(self._logits, (hidden_states.shape[0], self._logits.shape[-1]))
+        scores = empty_from(self._scores, (hidden_states.shape[0], self.n_top_k))
+        assignments = empty_from(self._assignments, (hidden_states.shape[0], self.n_top_k))
+        offsets = empty_from(self._offsets, (hidden_states.shape[0], self.n_top_k))
+        mapped_slots = empty_from(self._mapped_slots, (hidden_states.shape[0], self.n_top_k))
+        moe_input = empty_from(self._moe_input, (hidden_states.shape[0] * self.n_top_k, self._moe_input.shape[-1]))
+
+        self._gate_proj(logits, hidden_states, gate_w)
+        self._expert_counts.zero_()
+        self._top_1_gate(self._expert_counts, scores, assignments, offsets, logits, batch_metadata)
+        self._moe_scatter(moe_input, self._expert_cumsum, mapped_slots, hidden_states, self._expert_counts,
+                          assignments, offsets)
+
+        return moe_input, self._expert_cumsum, scores, mapped_slots
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                batch_metadata: RaggedBatchWrapper,
+                gate_w: torch.Tensor,
+                mlp_1_w: torch.Tensor,
+                mlp_2_w: torch.Tensor,
+                mlp_1_b: Optional[torch.Tensor] = None,
+                mlp_2_b: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        MoE forward pass built on top of CUTLASS multi-GEMM.
+
+        Parameters:
+            hidden_states (torch.Tensor): Hidden states tensor. Expected shape is [batch, seq_len, model_dim].
+            gate_w (torch.Tensor): Gate weight tensor. Expected shape is [num_experts, model_dim].
+        """
+
+        moe_input, expert_cumsum, scores, mapped_slots = self._gate(hidden_states, batch_metadata, gate_w)
+
+        # Get views on the buffers for GEMM
+        intermediate = empty_from(self._intermediate,
+                                  (hidden_states.shape[0] * self.n_top_k, self._intermediate.shape[-1]))
+        output_unordered = empty_from(self._output_unordered,
+                                      (hidden_states.shape[0] * self.n_top_k, self._output_unordered.shape[-1]))
+        output = empty_from(self._output, (hidden_states.shape[0], self._output.shape[-1]))
+
+        if self._activation is not None:
+            gated_intermediate = empty_from(
+                self._gated_intermediate, (hidden_states.shape[0] * self.n_top_k, self._gated_intermediate.shape[-1]))
+            self._mlp_1(
+                gated_intermediate,
+                moe_input,
+                mlp_1_w,
+                expert_cumsum,
+                mlp_1_b,
+            )
+            self._activation(intermediate, gated_intermediate)
+        else:
+            self._mlp_1(
+                intermediate,
+                moe_input,
+                mlp_1_w,
+                expert_cumsum,
+                mlp_1_b,
+            )
+
+        self._mlp_2(
+            output_unordered,
+            intermediate,
+            mlp_2_w,
+            expert_cumsum,
+            mlp_2_b,
+        )
+
+        self._moe_gather(output, output_unordered, scores, mapped_slots, self._expert_counts)
+        return output

deepspeed/inference/v2/modules/implementations/post_norm/__init__.py

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

deepspeed/inference/v2/modules/implementations/post_norm/cuda_post_ln.py

@@ -0,0 +1,56 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Tuple
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from ...interfaces import DSPostNormBase, DSPostNormRegistry
+from ...configs import DSNormConfig
+from ....kernels.core_ops.cuda_layer_norm.cuda_post_ln import CUDAFPPostLN
+from ....allocator import empty_from
+from ....inference_parameter import InferenceParameter
+
+
+@DSPostNormRegistry.register_module
+class DSPostLNCUDAModule(DSPostNormBase):
+
+    @staticmethod
+    def name():
+        return 'cuda_post_ln'
+
+    @staticmethod
+    def supports_config(config: DSNormConfig):
+        if len(set([config.residual_dtype, config.input_dtype, config.output_dtype])) != 1:
+            return False
+
+        try:
+            _ = CUDAFPPostLN(config.channels, config.residual_dtype)
+        except ValueError:
+            return False
+        return True
+
+    def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]):
+        super().__init__(config, implementation_config)
+        self._fp_post_ln = CUDAFPPostLN(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps)
+
+        self._output = torch.empty((config.max_tokens, config.channels),
+                                   dtype=config.output_dtype,
+                                   device=get_accelerator().current_device())
+
+    def transform_param(self, param: torch.Tensor) -> InferenceParameter:
+        param = param.to(self._config.input_dtype)
+        return InferenceParameter.initialize(param)
+
+    def forward(self, residual: torch.Tensor, hidden_in: torch.Tensor, gamma: torch.Tensor,
+                beta: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Since the CUDA FP only supports all data types being the same, we will alias the residual
+        with our output.
+        """
+        self._residual_output = empty_from(self._output, residual.shape)
+        self._fp_post_ln(residual, residual, hidden_in, gamma, beta)
+        return residual, residual

deepspeed/inference/v2/modules/implementations/pre_norm/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .cuda_pre_ln import DSPreLNCUDAModule
+from .cuda_pre_rms import DSPreRMSCUDAModule

deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_ln.py

@@ -0,0 +1,69 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from ...interfaces import DSPreNormBase, DSPreNormRegistry
+from ...configs import DSNormConfig, NormTypeEnum
+from ....kernels.core_ops.cuda_layer_norm.cuda_pre_ln import CUDAFPPreLN
+from ....kernels.core_ops.cuda_layer_norm.cuda_ln import CUDAFPLN
+from ....allocator import empty_from
+from ....inference_parameter import InferenceParameter
+
+
+@DSPreNormRegistry.register_module
+class DSPreLNCUDAModule(DSPreNormBase):
+
+    @staticmethod
+    def name():
+        return 'cuda_pre_ln'
+
+    @staticmethod
+    def supports_config(config: DSNormConfig):
+        type = NormTypeEnum(config.type)
+        if type != NormTypeEnum.LayerNorm:
+            return False
+
+        if len(set([config.residual_dtype, config.input_dtype, config.output_dtype])) != 1:
+            return False
+
+        try:
+            _ = CUDAFPPreLN(config.channels, config.residual_dtype)
+        except ValueError:
+            return False
+        return True
+
+    def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]):
+        super().__init__(config, implementation_config)
+        self._fp_pre_ln = CUDAFPPreLN(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps)
+        self._fp_ln = CUDAFPLN(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps)
+
+        # Buffers for the hidden output (residual is updated in-place)
+        self._hidden_output = torch.empty((config.max_tokens, config.channels),
+                                          dtype=config.output_dtype,
+                                          device=get_accelerator().current_device())
+
+    def transform_param(self, param: torch.Tensor) -> InferenceParameter:
+        param = param.to(self._config.input_dtype)
+        return InferenceParameter.initialize(param)
+
+    def forward(self, residual: torch.Tensor, hidden_in: Optional[torch.Tensor], gamma: torch.Tensor,
+                beta: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Since the CUDA FP only supports all data types being the same, we will alias the residual
+        with our output.
+
+        If hidden_in is None, that means we do not need to perform the residual add and will
+        only return the hidden output modified.
+        """
+        hidden_out = empty_from(self._hidden_output, residual.shape)
+        if hidden_in is None:
+            self._fp_ln(hidden_out, residual, gamma, beta)
+        else:
+            self._fp_pre_ln(residual, hidden_out, residual, hidden_in, gamma, beta)
+        return residual, hidden_out

deepspeed/inference/v2/modules/implementations/pre_norm/cuda_pre_rms.py

@@ -0,0 +1,79 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from ...interfaces import DSPreNormBase, DSPreNormRegistry
+from ...configs import DSNormConfig, NormTypeEnum
+from ....kernels.core_ops import CUDARMSNorm, CUDARMSPreNorm
+from ....allocator import empty_from
+from ....inference_parameter import InferenceParameter
+
+
+@DSPreNormRegistry.register_module
+class DSPreRMSCUDAModule(DSPreNormBase):
+
+    @staticmethod
+    def name():
+        return 'cuda_pre_rms'
+
+    @staticmethod
+    def supports_config(config: DSNormConfig):
+        type = NormTypeEnum(config.type)
+        if type != NormTypeEnum.RMSNorm:
+            return False
+
+        if len(set([config.residual_dtype, config.input_dtype, config.output_dtype])) != 1:
+            return False
+
+        try:
+            # Only need to check one since the support matrix for the two rms kernels is the same
+            _ = CUDARMSPreNorm(config.channels, config.residual_dtype)
+        except ValueError:
+            return False
+        return True
+
+    def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]):
+        super().__init__(config, implementation_config)
+        self._fp_rms = CUDARMSNorm(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps)
+        self._fp_rms_pre = CUDARMSPreNorm(self._config.channels, self._config.residual_dtype, epsilon=self._config.eps)
+
+        # Buffers for both the hidden and residual outputs
+        self._hidden_output = torch.empty((config.max_tokens, config.channels),
+                                          dtype=config.output_dtype,
+                                          device=get_accelerator().current_device())
+        self._residual_output = torch.empty((config.max_tokens, config.channels),
+                                            dtype=config.output_dtype,
+                                            device=get_accelerator().current_device())
+
+    def transform_param(self, param: torch.Tensor) -> InferenceParameter:
+        param = param.to(self._config.input_dtype)
+        return InferenceParameter.initialize(param)
+
+    def forward(self,
+                residual: torch.Tensor,
+                hidden_in: Optional[torch.Tensor],
+                gamma: torch.Tensor,
+                beta: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Since the CUDA FP only supports all data types being the same, we will alias the residual
+        with our output.
+
+        If hidden_in is None, that means we do not need to perform the residual add and will
+        only return the hidden output modified.
+        """
+        assert beta is None, "Beta is not supported for RMSNorm"
+
+        hidden_out = empty_from(self._hidden_output, residual.shape)
+        if hidden_in is None:
+            self._fp_rms(hidden_out, residual, gamma)
+            residual_out = residual
+        else:
+            residual_out = empty_from(self._residual_output, residual.shape)
+            self._fp_rms_pre(residual_out, hidden_out, residual, hidden_in, gamma)
+        return residual_out, hidden_out

deepspeed/inference/v2/modules/implementations/unembed/__init__.py

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

deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py

@@ -0,0 +1,120 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from ....allocator import empty_from
+from ....inference_utils import DtypeEnum, ActivationType
+from ....kernels.core_ops import CUDAFPLN, BlasLibLinear, CUDARMSNorm, CUDABiasActivation
+from ....kernels.ragged_ops import RaggedLogitsGather
+from ....ragged import RaggedBatchWrapper
+from ...interfaces import DSUnembedBase, DSUnembedRegistry
+from ...configs import DSUnembedConfig
+
+
+@DSUnembedRegistry.register_module
+class DSRaggedUnembed(DSUnembedBase):
+    """
+    Ragged unembedding implementation. This implementation will gather only the last token
+    of each sequence in the ragged inflight batch and calculate the logits only for those rows.
+    """
+
+    @staticmethod
+    def name():
+        return 'ragged_unembed'
+
+    @staticmethod
+    def supports_config(config: DSUnembedConfig):
+
+        if DtypeEnum(config.dtype) not in [DtypeEnum.fp16, DtypeEnum.bf16, DtypeEnum.fp32]:
+            return False
+
+        try:
+            _ = RaggedLogitsGather(config.model_dim, config.dtype)
+        except ValueError:
+            return False
+
+        if config.norm_type == 'rms_norm':
+            try:
+                _ = CUDARMSNorm(config.model_dim, config.dtype)
+            except ValueError:
+                return False
+        elif config.norm_type == 'layer_norm':
+            try:
+                _ = CUDAFPLN(config.model_dim, config.dtype)
+            except ValueError:
+                return False
+
+        return True
+
+    def __init__(self, config: DSUnembedConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+        self._logits_gather = RaggedLogitsGather(config.model_dim, self._config.dtype)
+
+        if self._config.norm_type == 'layer_norm':
+            self._norm = CUDAFPLN(self._config.model_dim, self._config.dtype)
+        elif self._config.norm_type == 'rms_norm':
+            self._norm = CUDARMSNorm(self._config.model_dim, self._config.dtype)
+        else:
+            self._norm = None
+
+        self._linear = BlasLibLinear(self._config.dtype)
+        # Here the activation kernel is being used to apply bias, hence the identity activation type!
+        self._act_fn = CUDABiasActivation(self._config.vocab_size, self._config.dtype, ActivationType.IDENTITY)
+
+        self._intermediate = torch.empty((self._config.max_sequences, self._config.model_dim),
+                                         dtype=self._config.dtype,
+                                         device=get_accelerator().current_device())
+
+        self._output = torch.empty((self._config.max_sequences, self._config.vocab_size),
+                                   dtype=self._config.dtype,
+                                   device=get_accelerator().current_device())
+
+    @property
+    def output(self) -> torch.Tensor:
+        return self._output
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                vocab_embedding: torch.Tensor,
+                ragged_metadata: RaggedBatchWrapper,
+                bias: Optional[torch.Tensor] = None,
+                gamma: Optional[torch.Tensor] = None,
+                beta: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Return final model logits.
+
+        Args:
+            hidden_states (torch.Tensor): The hidden states from the model. This is the output of the
+                final layer of the model.
+            vocab_embedding (torch.Tensor): The vocab embedding table.
+            raged_metadata (RaggedBatchWrapper): The ragged batch metadata.
+            gamma (Optional[torch.Tensor]): The gamma tensor for normalization.
+            beta (Optional[torch.Tensor]): The beta tensor for normalization.
+        """
+
+        cut_down_hidden_states = empty_from(self._intermediate,
+                                            (ragged_metadata.current_sequences, self._config.model_dim))
+        self._logits_gather(cut_down_hidden_states, hidden_states, ragged_metadata)
+
+        if self._config.norm_type == 'rms_norm':
+            if gamma is None:
+                raise ValueError('RMS Normalization enabled but gamma not provided.')
+            self._norm(cut_down_hidden_states, cut_down_hidden_states, gamma)
+        elif self._config.norm_type == 'layer_norm':
+            if gamma is None or beta is None:
+                raise ValueError('Normalization enabled but gamma and/or beta not provided.')
+            self._norm(cut_down_hidden_states, cut_down_hidden_states, gamma, beta)
+
+        output = empty_from(self._output, (ragged_metadata.current_sequences, self._config.vocab_size))
+        self._linear(output, cut_down_hidden_states, vocab_embedding)
+        if bias is not None:
+            self._act_fn(output, bias)
+
+        return output

deepspeed/inference/v2/modules/interfaces/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from .attention_base import DSSelfAttentionRegistry, DSSelfAttentionBase
+from .embedding_base import DSEmbeddingRegistry, DSEmbeddingBase
+from .linear_base import DSLinearRegistry, DSLinearBase
+from .moe_base import DSMoERegistry, DSMoEBase
+from .post_norm_base import DSPostNormRegistry, DSPostNormBase
+from .pre_norm_base import DSPreNormRegistry, DSPreNormBase
+from .unembed_base import DSUnembedRegistry, DSUnembedBase

deepspeed/inference/v2/modules/interfaces/attention_base.py

@@ -0,0 +1,97 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional, Tuple, Type
+
+import torch
+
+from ...ragged import RaggedBatchWrapper
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ..ds_module import DSModuleBase
+from ..module_registry import DSModuleRegistryBase
+from ..configs import DSSelfAttentionConfig
+
+
+class DSSelfAttentionBase(DSModuleBase):
+    """
+    Base mixin for all attention modules. The interface represented by this module
+    is broadly:
+
+    output = attention(query_key_value,
+                       Optional[kv_cache],
+                       Optional[attention_mask],
+                       Optional[attention_bias])
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSSelfAttentionConfig
+
+    def __init__(self, config: DSSelfAttentionConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+    @property
+    def kv_block_size(self) -> int:
+        """
+        Return preferred granulatity for blocked KV-cache implementation.
+        """
+        raise NotImplementedError()
+
+    @property
+    def q_block_size(self) -> int:
+        """
+        Property to calculate blocking granularity for the query dimension.
+        This has no impact on the KV-cache structure, but will  affect the
+        number of attention atoms associated with a batch.
+        """
+        raise NotImplementedError()
+
+    def build_atoms(self, ragged_batch: RaggedBatchWrapper) -> None:
+        """
+        Build the atoms for this module. This is not a strict requirement for the class,
+        so this method is a no-op by default rather than abstract.
+        """
+        pass
+
+    def forward(self,
+                q_k_v: torch.Tensor,
+                kv_cache: torch.Tensor,
+                batch: RaggedBatchWrapper,
+                attention_mask: Optional[torch.Tensor] = None,
+                attention_bias: Optional[torch.Tensor] = None,
+                inv_freqs: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Parameters:
+            q_k_v (torch.Tensor): Query, key, and value tensors. Expected shape is:
+                [
+                    batch,
+                    seq_len,
+                    2 * self._config.n_heads_kv + self._config.n_heads_q,
+                    self._config.head_size
+                ].
+            kv_cache (Optional[torch.Tensor]): Key and value cache tensor. Expected shape is
+                [
+                    2,
+                    batch,
+                    kv_cache_len,
+                    self._config.n_heads_kv,
+                    self._config.head_size
+                ]. If None, cache is disabled. The `kv_cache_len` dimension does not need to
+                be contiguous (it should expand stride by `max_out_tokens`).
+            batch (RaggedBatchWrapper): Ragged batch metadata.
+            attention_mask (Optional[torch.Tensor]): Attention mask tensor. If None, masking is
+                disabled. This will defer to the config in the case of conflicting information.
+                This means if the config class is implying causal attention, the mask will be ignored.
+            attention_bias (Optional[torch.Tensor]): Attention bias tensor. If None, bias is disabled.
+        """
+        raise NotImplementedError()
+
+
+class DSSelfAttentionRegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSSelfAttentionBase

deepspeed/inference/v2/modules/interfaces/embedding_base.py

@@ -0,0 +1,85 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import abstractmethod
+from typing import Any, Dict, Optional, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ...ragged import RaggedBatchWrapper
+from ..ds_module import DSModuleBase
+from ..module_registry import DSModuleRegistryBase
+from ..configs import DSEmbeddingsConfig
+from ...inference_parameter import InferenceParameter
+
+
+class DSEmbeddingBase(DSModuleBase):
+    """
+    Base mixin for embedding modules. The interface represented by this module is:
+
+    hidden_out = embedding(input_ids) +
+                 position_embedding(position_ids) +
+                 token_type_embedding(token_type_ids)
+    with optional normalization.
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSEmbeddingsConfig
+
+    def __init__(self, config: DSEmbeddingsConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+    def transform_param(self, embed_param: torch.Tensor) -> InferenceParameter:
+        """
+        Perform any necessary transformations on an embedding parameter. This module assumes
+        that all embedding parameters would require the same set of transformations.
+
+        Parameters:
+            embed_param (torch.Tensor): Embedding parameter. Shape is of [vocab_size, hidden_size]
+        """
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def output(self) -> torch.Tensor:
+        """
+        Pre-allocated output Tensor. This currently needs to be exposed for gather operations
+        on the output.
+
+        TODO(cmikeh2): This is not ideal. We need a better abstraction for this, such as giving
+        access to the inference comm object to the DSModule.
+        """
+        raise NotImplementedError()
+
+    def forward(self,
+                ragged_batch: RaggedBatchWrapper,
+                word_embeddings: torch.Tensor,
+                position_embeddings: Optional[torch.Tensor] = None,
+                token_type_ids: Optional[torch.Tensor] = None,
+                token_type_embeddings: Optional[torch.Tensor] = None) -> InferenceParameter:
+        """
+        Parameters:
+            ragged_batch (torch.Tensor): Ragged batch of token ids + associated metadata.
+            word_embeddings (torch.Tensor): Word embeddings.
+            position_embeddings (torch.Tensor): Position embeddings. If passed, IDs will be
+                inferred from the ragged batch itself.
+            token_type_ids (torch.Tensor): Token type ids.
+            token_type_embeddings (torch.Tensor): Token type embeddings.
+
+        Returns:
+            torch.Tensor: Hidden states. This should be the sum of the relevant
+                encodings for the model.
+        """
+        raise NotImplementedError()
+
+
+class DSEmbeddingRegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSEmbeddingBase

deepspeed/inference/v2/modules/interfaces/linear_base.py

@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import abstractmethod
+from typing import Any, Dict, Optional, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ..ds_module import DSModuleBase
+from ..module_registry import DSModuleRegistryBase
+from ..configs import DSLinearConfig
+from ...inference_parameter import InferenceParameter
+
+
+class DSLinearBase(DSModuleBase):
+    """
+    Base mixin for all Linear modules. The interface represented by this module
+    is:
+
+    hidden_out = activation(hidden_in * weight + bias)
+
+    The format and dtype of the weight and bias tensors are not defined and implementations
+    may compress as necessary. Must support a bias.
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSLinearConfig
+
+    def __init__(self, config: DSLinearConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+    @abstractmethod
+    def transform_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Perform any necessary transformations of the parameters of this module.
+
+        Parameters:
+            param (torch.Tensor): Weight or bias tensor.
+        """
+        ...
+
+    def forward(self, hidden_states: torch.Tensor, w: torch.Tensor, b: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Parameters:
+            hidden_states (torch.Tensor): Hidden states tensor. Expected shape is either
+                [batch, seq_len, in_channels] or [batch, in_channels].
+
+        Returns:
+            torch.Tensor: Output tensor. Tensor should have same number of dimensions as
+                input tensor.
+        """
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def output(self) -> torch.Tensor:
+        """
+        Return the padded, pre-allocated output Tensor.
+        """
+        ...
+
+
+class DSLinearRegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSLinearBase

deepspeed/inference/v2/modules/interfaces/moe_base.py

@@ -0,0 +1,91 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import abstractmethod
+from typing import Any, Dict, Optional, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ..ds_module import DSModuleBase
+from ..module_registry import DSModuleRegistryBase
+from ..configs import DSMoEConfig
+from ...inference_parameter import InferenceParameter
+
+
+class DSMoEBase(DSModuleBase):
+    """
+    Base mixing for MoE modules. The interface represented by this module is:
+
+    expert_assignments = gate(hidden_states)
+    intermediate = ragged_linear(hidden_states, expert_assignments)
+    output = ragged_linear(intermediate, expert_assignments)
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSMoEConfig
+
+    def __init__(self, config: DSMoEConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+    @abstractmethod
+    def transform_gate_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Perform any necessary transformations of the gate parameter.
+
+        Args:
+            param (torch.Tensor): gate_w (shape: [num_experts, model_dim])
+        """
+        ...
+
+    @abstractmethod
+    def transform_moe_mlp_1_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Perform any necessary transformations of the parameter. The specific component
+        being transformed should be inferred from the shape of the parameter.
+
+        Args:
+            param (torch.Tensor): One of either mlp_1_w, mlp_1_b
+        """
+        ...
+
+    @abstractmethod
+    def transform_moe_mlp_2_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Perform any necessary transformations of the parameter. The specified component being
+        transformed should be inferred from the shape of the parameter. This interface is
+        separate from transform_moe_1_param because the two components may have identical
+        shapes.
+
+        Args:
+            param (torch.Tensor): One of either mlp_2_w or mlp_2_b
+        """
+        ...
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                gate_w: torch.Tensor,
+                mlp_1_w: torch.Tensor,
+                mlp_2_w: torch.Tensor,
+                mlp_1_b: Optional[torch.Tensor] = None,
+                mlp_2_b: Optional[torch.Tensor] = None) -> torch.Tensor:
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def output(self) -> torch.Tensor:
+        """
+        Returns the pre-allocated, padded output Tensor.
+        """
+        ...
+
+
+class DSMoERegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSMoEBase

deepspeed/inference/v2/modules/interfaces/post_norm_base.py

@@ -0,0 +1,69 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import abstractmethod
+from typing import Any, Dict, Optional, Tuple, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ..ds_module import DSModuleBase
+from ..configs.norm_config import DSNormConfig
+from ..module_registry import DSModuleRegistryBase
+from ...inference_parameter import InferenceParameter
+
+
+class DSPostNormBase(DSModuleBase):
+    """
+    Base MixIn for all Post-Normalization modules. The interface represented by this
+    module is:
+
+    residual, hidden_out = norm(residual + hidden_in)
+
+    If residual and hidden_out are the same data type, then they may alias each other.
+    Furthermore, residual should be updated in-place.
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSNormConfig
+
+    def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+    @abstractmethod
+    def transform_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Transform a gamma/beta parameter. It is assumed that both transformations are
+        the same.
+
+        Parameters:
+            param (torch.Tensor): Gamma or beta parameter.
+        """
+        ...
+
+    def forward(self,
+                residual: torch.Tensor,
+                hidden_states: torch.Tensor,
+                gamma: torch.Tensor,
+                beta: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Parameters:
+            residual (torch.Tensor): Residual tensor.
+            hidden_states (torch.Tensor): Hidden states tensor.
+
+        Returns:
+            (torch.Tensor, torch.Tensor): Tuple of residual and hidden states.
+                Hidden states may alias with residual.
+        """
+        raise NotImplementedError()
+
+
+class DSPostNormRegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSPostNormBase

deepspeed/inference/v2/modules/interfaces/pre_norm_base.py

@@ -0,0 +1,73 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import abstractmethod
+from typing import Any, Dict, Optional, Tuple, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ..ds_module import DSModuleBase
+from ..configs.norm_config import DSNormConfig
+from ..module_registry import DSModuleRegistryBase
+from ...inference_parameter import InferenceParameter
+
+
+class DSPreNormBase(DSModuleBase):
+    """
+    Base mixin for all Pre-Normalization modules. The interface represented by this module
+    is:
+
+    if hidden_in is not None:
+        residual_out = residual + hidden_in
+    else:
+        residual_out = residual
+
+    hidden_out = normalize(residual_out)
+    return residual_out, hidden_out
+
+    Residual should be updated in-place.
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSNormConfig
+
+    def __init__(self, config: DSNormConfig, implementation_config: Dict[str, Any]):
+        super().__init__(config, implementation_config)
+
+    @abstractmethod
+    def transform_param(self, param: torch.Tensor) -> InferenceParameter:
+        """
+        Transform a gamma/beta parameter. It is assumed that both transformations are
+        the same.
+
+        Parameters:
+            param (torch.Tensor): Gamma or beta parameter.
+        """
+        ...
+
+    def forward(self,
+                residual: torch.Tensor,
+                hidden_states: Optional[torch.Tensor],
+                gamma: torch.Tensor,
+                beta: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Parameters:
+            residual (torch.Tensor): Residual tensor.
+            hidden_states (torch.Tensor): Hidden states tensor.
+
+        Returns:
+            (torch.Tensor, torch.Tensor): Tuple of residual and hidden states.
+        """
+        raise NotImplementedError()
+
+
+class DSPreNormRegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSPreNormBase

deepspeed/inference/v2/modules/interfaces/unembed_base.py

@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Any, Dict, Optional, Type
+
+import torch
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from ...ragged import RaggedBatchWrapper
+from ..ds_module import DSModuleBase
+from ..module_registry import DSModuleRegistryBase
+from ..configs import DSUnembedConfig
+
+
+class DSUnembedBase(DSModuleBase):
+    """
+    Base mixin for unmebedding modules. The interface represented by this module is:
+
+    if config.do_normalization
+        hidden = layer_norm(hidden)
+    logits = hidden @ projection
+    """
+
+    @staticmethod
+    def config_class() -> Type[DeepSpeedConfigModel]:
+        return DSUnembedConfig
+
+    def __init__(self, config: DSUnembedConfig, implementation_config: Dict[str, Any]) -> None:
+        super().__init__(config, implementation_config)
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                vocab_embedding: torch.Tensor,
+                ragged_metadata: RaggedBatchWrapper,
+                gamma: Optional[torch.Tensor] = None,
+                beta: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Forward interface. Gamma and beta are optional parameters passed depending on
+        `self.config.do_normalization`.
+
+        Args:
+            hidden_states (torch.Tensor): Hidden states of shape [tokens, model_dim]
+            vocab_embedding (torch.Tensor): Embedding matrix of shape [vocab_size, model_dim]
+            ragged_metadata (RaggedBatchWrapper): Metadata for the ragged batch.
+            gamma (Optional[torch.Tensor]): Gamma parameter for layer norm.
+            beta (Optional[torch.Tensor]): Beta parameter for layer norm.
+
+        Returns:
+            torch.Tensor: Unembedded hidden states of shape [n_seqs, model_dim]
+        """
+        raise NotImplementedError()
+
+
+class DSUnembedRegistry(DSModuleRegistryBase):
+    registry: Dict = {}
+
+    @staticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        return DSUnembedBase

deepspeed/inference/v2/modules/module_registry.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from abc import ABC, abstractstaticmethod
+from typing import Any, Dict, Type
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+from .ds_module import DSModuleBase
+
+
+class ConfigBundle(DeepSpeedConfigModel):
+    """
+    A config bundle is a collection of configs that are used to instantiate a model implementation.
+    """
+    name: str
+    config: DeepSpeedConfigModel
+    implementation_config: Dict[str, Any] = {}
+
+
+class DSModuleRegistryBase(ABC):
+    """
+    Class holding logic for tracking the DSModule implementations of a given interface.
+    """
+
+    @classmethod
+    def instantiate_config(cls, config_bundle: ConfigBundle) -> DSModuleBase:
+        """
+        Given a DSModule key, attempt to instantiate
+        """
+        if config_bundle.name not in cls.registry:
+            raise KeyError(f"Unknown DSModule: {config_bundle.name}, cls.registry={cls.registry}")
+
+        target_implementation = cls.registry[config_bundle.name]
+        if not target_implementation.supports_config(config_bundle.config):
+            raise ValueError(f"Config {config_bundle.config} is not supported by {target_implementation}")
+
+        return cls.registry[config_bundle.name](config_bundle.config, config_bundle.implementation_config)
+
+    @abstractstaticmethod
+    def associated_class() -> Type[DSModuleBase]:
+        """
+        Return the class associated with this registry.
+        """
+        raise NotImplementedError("Must associated a DSModule class with its registry.")
+
+    @classmethod
+    def register_module(cls, child_class: DSModuleBase) -> None:
+        """
+        Register a module with this registry.
+        """
+        if not issubclass(child_class, cls.associated_class()):
+            raise TypeError(
+                f"Can only register subclasses of {cls.associated_class()}, {child_class} does not inherit from {cls.associated_class()}"
+            )
+        cls.registry[child_class.name()] = child_class
+        return child_class