chore: import upstream snapshot with attribution

tests/unit/inference/__init__.py

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

tests/unit/inference/quantization/test_intX_quantization.py

@@ -0,0 +1,419 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import numpy as np
+import torch
+import torch.nn as nn
+from unit.common import DistributedTest
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.quantization.quantization import _init_group_wise_weight_quantization
+from deepspeed.inference.quantization.utils import Quantizer, DeQuantizer
+from deepspeed.inference.quantization.layers import QuantizedLinear
+from deepspeed.utils.torch import required_torch_version
+from transformers.models.opt.modeling_opt import OPTDecoderLayer
+from transformers import AutoConfig, OPTConfig, AutoModel
+import pytest
+from collections import OrderedDict
+from typing import Dict
+from deepspeed.ops.aio import AsyncIOBuilder
+
+device = get_accelerator().device_name() if get_accelerator().is_available() else 'cpu'
+
+if not required_torch_version(min_version=1.11):
+    pytest.skip("torch.Tensor.bitwise_left_shift in INT4 quantizer needs torch 1.11 or above.",
+                allow_module_level=True)
+
+
+def reset_random(seed=1234):
+    np.random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    get_accelerator().manual_seed_all(seed)
+
+
+def quantization_test_helper(pre_quant_type: torch.dtype, num_bits: int):
+    reset_random()
+    num_group = 1024 * 32
+    group_size = 64
+    quantization_config = {'num_bits': num_bits, 'group_size': group_size, 'group_dim': 1, 'symmetric': False}
+
+    quantizer = Quantizer(config=quantization_config)
+    dequantizer = DeQuantizer(config=quantization_config, dtype=pre_quant_type)
+
+    data = torch.randn(num_group, group_size, dtype=pre_quant_type, device=device)
+
+    quantized_data, scale_buf, min_vals = quantizer.quantize(data)
+    dequantized_data = dequantizer.dequantize(quantized_data, scale_buf, min_vals)
+
+    max_diff = torch.max(torch.abs(data - dequantized_data))
+    mean_diff = torch.mean(torch.abs(data - dequantized_data))
+
+    # This threshold value is emperically selected.
+    assert mean_diff < 0.15 and max_diff < 0.5, f'Numeric error exceed threshold, mean diff {mean_diff} (threshold 0.15), max diff {max_diff} (threshold 0.5)'
+
+
+def zero3_post_init_quantization_test_helper(cpu_offload: bool, nvme_offload: bool, bits: int, nvme_path=None):
+    import deepspeed
+    from transformers.integrations.deepspeed import HfDeepSpeedConfig
+
+    if nvme_offload and not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]:
+        pytest.skip('Skip tests since async-io is not compatible')
+
+    def get_zero3_ds_config(hf_config: OPTConfig, cpu_offload: bool, nvme_offload: bool, bits: int) -> Dict:
+        GB = 1 << 30
+
+        ds_config = {
+            "fp16": {
+                "enabled": True,
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "stage3_prefetch_bucket_size": 2 * hf_config.hidden_size * hf_config.hidden_size,
+                "stage3_param_persistence_threshold": hf_config.hidden_size,
+                "stage3_max_live_parameters": 2 * hf_config.hidden_size * hf_config.hidden_size
+            },
+            "steps_per_print": 2000,
+            "train_micro_batch_size_per_gpu": 1,
+            "wall_clock_breakdown": False,
+            'weight_quantization': {
+                'post_init_quant': {
+                    'fc': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                    'self_attn.q_proj': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                    'self_attn.k_proj': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                    'self_attn.v_proj': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                    'self_attn.out_proj': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                    'lm_head': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                    'embed_tokens': {
+                        'num_bits': bits,
+                        'group_size': 32,
+                        'group_dim': 1,
+                        'symmetric': False
+                    },
+                }
+            }
+        }
+
+        if cpu_offload:
+            ds_config["zero_optimization"]["offload_param"] = dict(device="cpu", pin_memory=1)
+        if nvme_offload:
+            ds_config["zero_optimization"]["offload_param"] = dict(
+                device="nvme",
+                pin_memory=True,
+                nvme_path=nvme_path or '~/tmp_offload_dir',
+                buffer_count=5,
+                buffer_size=1 * GB,
+            )
+            ds_config["aio"] = {
+                "block_size": 1048576,
+                "queue_depth": 8,
+                "thread_count": 1,
+                "single_submit": False,
+                "overlap_events": True,
+            }
+
+        return ds_config
+
+    hf_config = AutoConfig.from_pretrained('facebook/opt-125m')
+    ds_config = get_zero3_ds_config(hf_config=hf_config, cpu_offload=cpu_offload, nvme_offload=nvme_offload, bits=bits)
+
+    input_ids = torch.ones(1, 16, dtype=torch.int32, device=device)
+    attention_mask = torch.ones(1, 16, dtype=torch.float32, device=device)
+
+    with torch.no_grad():
+        ref_model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16).to(device)
+        ref_model.eval()
+        ref_output = ref_model(input_ids=input_ids, attention_mask=attention_mask)
+
+    with torch.no_grad():
+        dschf = HfDeepSpeedConfig(ds_config)
+        model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16)
+        model = model.eval()
+
+        model = _init_group_wise_weight_quantization(model, ds_config)
+        ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
+        ds_engine.module.eval()
+        model = ds_engine.module
+
+    output = model(input_ids=input_ids, attention_mask=attention_mask)
+
+    mean_diff = torch.mean(torch.abs(output.last_hidden_state - ref_output.last_hidden_state))
+
+    # This threshold value is emperically selected.
+    assert mean_diff < 0.4, f'Numeric error exceed threshold, relative error {mean_diff} (threshold 0.4)'
+
+
+def zero3_quantized_initialization_test_helper(cpu_offload: bool, nvme_offload: bool, bits: int, nvme_path=None):
+    import deepspeed
+    from transformers.integrations.deepspeed import HfDeepSpeedConfig
+
+    if nvme_offload and not deepspeed.ops.__compatible_ops__[AsyncIOBuilder.NAME]:
+        pytest.skip('Skip tests since async-io is not compatible')
+
+    def get_zero3_ds_config(hf_config: OPTConfig, cpu_offload: bool, nvme_offload: bool, bits: int) -> Dict:
+        GB = 1 << 30
+
+        ds_config = {
+            "fp16": {
+                "enabled": True,
+            },
+            "zero_optimization": {
+                "stage": 3,
+                "stage3_prefetch_bucket_size": 2 * hf_config.hidden_size * hf_config.hidden_size,
+                "stage3_param_persistence_threshold": hf_config.hidden_size,
+                "stage3_max_live_parameters": 2 * hf_config.hidden_size * hf_config.hidden_size
+            },
+            "steps_per_print": 2000,
+            "train_micro_batch_size_per_gpu": 1,
+            "wall_clock_breakdown": False,
+            'weight_quantization': {
+                'quantized_initialization': {
+                    'num_bits': bits,
+                    'group_size': 32,
+                    'group_dim': 1,
+                    'symmetric': False
+                },
+            }
+        }
+
+        if cpu_offload:
+            ds_config["zero_optimization"]["offload_param"] = dict(device="cpu", pin_memory=1)
+        if nvme_offload:
+            ds_config["zero_optimization"]["offload_param"] = dict(
+                device="nvme",
+                pin_memory=True,
+                nvme_path=nvme_path or '~/tmp_offload_dir',
+                buffer_count=5,
+                buffer_size=1 * GB,
+            )
+            ds_config["aio"] = {
+                "block_size": 1048576,
+                "queue_depth": 8,
+                "thread_count": 1,
+                "single_submit": False,
+                "overlap_events": True,
+            }
+
+        return ds_config
+
+    hf_config = AutoConfig.from_pretrained('facebook/opt-125m')
+    ds_config = get_zero3_ds_config(hf_config=hf_config, cpu_offload=cpu_offload, nvme_offload=nvme_offload, bits=bits)
+
+    input_ids = torch.ones(1, 16, dtype=torch.int32, device=device)
+    attention_mask = torch.ones(1, 16, dtype=torch.float32, device=device)
+
+    with torch.no_grad():
+        ref_model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16).to(device)
+        ref_model.eval()
+        ref_output = ref_model(input_ids=input_ids, attention_mask=attention_mask)
+
+    with torch.no_grad():
+        dschf = HfDeepSpeedConfig(ds_config)
+        model = AutoModel.from_pretrained('facebook/opt-125m', torch_dtype=torch.float16)
+        model = model.eval()
+        ds_engine = deepspeed.initialize(model=model, config_params=ds_config)[0]
+        ds_engine.module.eval()
+        model = ds_engine.module
+
+    output = model(input_ids=input_ids, attention_mask=attention_mask)
+
+    mean_diff = torch.mean(torch.abs(output.last_hidden_state - ref_output.last_hidden_state))
+
+    # This threshold value is emperically selected.
+    assert mean_diff < 0.4, f'Numeric error exceed threshold, relative error {mean_diff} (threshold 0.4)'
+
+
+@pytest.fixture(params=[4, 8], ids=["4bits", "8bits"])
+def quantization_bits(request):
+    return request.param
+
+
+@pytest.fixture(params=[0, 1], ids=["0", "1"])
+def group_dim(request):
+    return request.param
+
+
+class TestQuantizedInt(DistributedTest):
+
+    def test_model_quantization(self, quantization_bits):
+        reset_random()
+
+        config = AutoConfig.from_pretrained('facebook/opt-125m')
+
+        with torch.no_grad():
+            model = OPTDecoderLayer(config).half().to(device)
+            bits = quantization_bits
+
+            ds_config = {
+                'weight_quantization': {
+                    'post_init_quant': {
+                        'fc': {
+                            'num_bits': bits,
+                            'group_size': 64,
+                            'group_dim': 0,
+                            'symmetric': False
+                        },
+                        'self_attn.q_proj': {
+                            'num_bits': bits,
+                            'group_size': 64,
+                            'group_dim': 0,
+                            'symmetric': False
+                        },
+                        'self_attn.k_proj': {
+                            'num_bits': bits,
+                            'group_size': 64,
+                            'group_dim': 0,
+                            'symmetric': False
+                        },
+                        'self_attn.v_proj': {
+                            'num_bits': bits,
+                            'group_size': 64,
+                            'group_dim': 0,
+                            'symmetric': False
+                        },
+                        'self_attn.out_proj': {
+                            'num_bits': bits,
+                            'group_size': 64,
+                            'group_dim': 0,
+                            'symmetric': False
+                        }
+                    }
+                }
+            }
+
+            model = _init_group_wise_weight_quantization(model, ds_config)
+
+            assert type(model.fc1) is QuantizedLinear
+            assert type(model.fc2) is QuantizedLinear
+            assert type(model.self_attn.q_proj) is QuantizedLinear
+            assert type(model.self_attn.k_proj) is QuantizedLinear
+            assert type(model.self_attn.v_proj) is QuantizedLinear
+            assert type(model.self_attn.out_proj) is QuantizedLinear
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_quantized_linear(self, quantization_bits, group_dim):
+        reset_random()
+
+        layers = []
+
+        for idx in range(5):
+            layers.append(
+                (f'layer_{idx}', nn.Linear(in_features=128, out_features=128, dtype=torch.float16, device=device)))
+
+        input_tensor = torch.randn(32, 128, dtype=torch.float16, device=device)
+        with torch.no_grad():
+            model = nn.Sequential(OrderedDict(layers))
+
+            ref_output = model(input_tensor)
+
+            ds_config = {
+                'weight_quantization': {
+                    'post_init_quant': {
+                        'layer': {
+                            'num_bits': quantization_bits,
+                            'group_size': 64,
+                            'group_dim': group_dim,
+                            'symmetric': False
+                        }
+                    }
+                }
+            }
+
+            model = _init_group_wise_weight_quantization(model, ds_config)
+
+            assert type(model.layer_0) is QuantizedLinear
+            assert type(model.layer_1) is QuantizedLinear
+            assert type(model.layer_2) is QuantizedLinear
+            assert type(model.layer_3) is QuantizedLinear
+            assert type(model.layer_4) is QuantizedLinear
+
+            output = model(input_tensor)
+
+            mean_diff = torch.mean(torch.abs(ref_output - output))
+
+            # This threshold value is emperically selected.
+            assert mean_diff < 0.15, f'Numeric error exceed threshold, mean diff {mean_diff}'
+
+    def test_float_int4_quantization(self):
+        reset_random()
+        quantization_test_helper(torch.float32, 4)
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_half_int4_quantization(self):
+        reset_random()
+        quantization_test_helper(torch.float16, 4)
+
+    def test_float_int8_quantization(self):
+        reset_random()
+        quantization_test_helper(torch.float32, 8)
+
+    def test_half_int8_quantization(self):
+        reset_random()
+        quantization_test_helper(torch.float16, 8)
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_zero3_int4_post_init_quant(self, quantization_bits):
+        reset_random()
+        zero3_post_init_quantization_test_helper(cpu_offload=False, nvme_offload=False, bits=quantization_bits)
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_zero3_int4_post_init_quant_cpu_offload(self, quantization_bits):
+        reset_random()
+        zero3_post_init_quantization_test_helper(cpu_offload=True, nvme_offload=False, bits=quantization_bits)
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_zero3_int4_post_init_quant_nvme_offload(self, tmpdir):
+        reset_random()
+        zero3_post_init_quantization_test_helper(cpu_offload=False,
+                                                 nvme_offload=True,
+                                                 bits=4,
+                                                 nvme_path=str(tmpdir.join("nvme_offload")))
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_zero3_int4_quantized_initialization(self, quantization_bits):
+        reset_random()
+        zero3_quantized_initialization_test_helper(cpu_offload=False, nvme_offload=False, bits=quantization_bits)
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_zero3_int4_quantized_initialization_cpu_offload(self, quantization_bits):
+        reset_random()
+        zero3_quantized_initialization_test_helper(cpu_offload=True, nvme_offload=False, bits=quantization_bits)
+
+    @pytest.mark.skipif(device == 'cpu', reason='CPU does support FP16 GEMM')
+    def test_zero3_int4_quantized_initialization_nvme_offload(self, tmpdir):
+        reset_random()
+        zero3_quantized_initialization_test_helper(cpu_offload=False,
+                                                   nvme_offload=True,
+                                                   bits=4,
+                                                   nvme_path=str(tmpdir.join("nvme_offload")))

tests/unit/inference/test_checkpoint_sharding.py

@@ -0,0 +1,150 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import pytest
+import torch
+import deepspeed
+from deepspeed.model_implementations import DeepSpeedTransformerInference
+from unit.common import DistributedTest, DistributedFixture
+from transformers import AutoConfig, AutoModelForCausalLM
+import deepspeed.comm as dist
+from huggingface_hub import snapshot_download
+from deepspeed.ops.op_builder import InferenceBuilder
+
+# Handle different versions of transformers
+try:
+    from transformers.utils import is_offline_mode
+except ImportError:
+    # For transformers >= 5.0, is_offline_mode was removed
+    # transformers >= 5.0 requires huggingface_hub >= 1.2.1 which has is_offline_mode
+    from huggingface_hub import is_offline_mode
+
+from deepspeed.accelerator import get_accelerator
+
+if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
+    pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+
+def check_dtype(model, expected_dtype):
+
+    def find_dtype(module):
+        for child in module.children():
+            if isinstance(child, DeepSpeedTransformerInference):
+                return child.attention.attn_qkvw.dtype
+            else:
+                found_dtype = find_dtype(child)
+                if found_dtype:
+                    return found_dtype
+
+    found_dtype = find_dtype(model)
+    assert found_dtype, "Did not find DeepSpeedTransformerInference in model"
+    assert (found_dtype == expected_dtype), f"Expected transformer dtype {expected_dtype}, but found {found_dtype}"
+
+
+@pytest.fixture(params=[
+    "bigscience/bloom-560m", "EleutherAI/gpt-j-6B", "EleutherAI/gpt-neo-125M", "facebook/opt-350m", "facebook/opt-125m"
+])
+def model_name(request):
+    return request.param
+
+
+@pytest.fixture(params=[torch.float16, torch.int8], ids=["fp16", "int8"])
+def dtype(request):
+    if request.param not in get_accelerator().supported_dtypes():
+        pytest.skip(f"{request.param} not supported by {get_accelerator().device_name()}.")
+    return request.param
+
+
+class save_shard(DistributedFixture):
+    world_size = 2
+
+    def run(self, model_name, class_tmpdir):
+        # Only write a checkpoint if one does not exist
+        if not os.path.isdir(os.path.join(class_tmpdir, model_name)):
+            world_size = int(os.getenv("WORLD_SIZE", "1"))
+            inf_config = {
+                "replace_with_kernel_inject": True,
+                "dtype": torch.float16,
+                "enable_cuda_graph": False,
+                "tensor_parallel": {
+                    "tp_size": world_size
+                },
+                "save_mp_checkpoint_path": os.path.join(str(class_tmpdir), model_name),
+            }
+
+            # Load model and save sharded checkpoint
+            model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
+            model = deepspeed.init_inference(model, config=inf_config)
+
+
+@pytest.mark.seq_inference
+class TestCheckpointShard(DistributedTest):
+    world_size = 2
+
+    def test(self, model_name, dtype, class_tmpdir, save_shard):
+        world_size = int(os.getenv("WORLD_SIZE", "1"))
+        inf_config = {
+            "replace_with_kernel_inject": True,
+            "dtype": dtype,
+            "enable_cuda_graph": False,
+            "tensor_parallel": {
+                "tp_size": world_size
+            },
+            "checkpoint": os.path.join(class_tmpdir, model_name, "ds_inference_config.json"),
+        }
+
+        # Load model on meta tensors
+        model_config = AutoConfig.from_pretrained(model_name)
+        # Note that we use half precision to load initially, even for int8
+        with deepspeed.OnDevice(dtype=torch.float16, device="meta"):
+            model = AutoModelForCausalLM.from_config(model_config, torch_dtype=torch.bfloat16)
+        model = model.eval()
+        model = deepspeed.init_inference(model, config=inf_config)
+        check_dtype(model, dtype)
+
+
+@pytest.mark.seq_inference
+class TestCheckpointShardinAutoTP(DistributedTest):
+    world_size = 2
+
+    def test(self, model_name, class_tmpdir):
+
+        def write_checkpoints_json(model_name, class_tmpdir):
+            import json
+            from pathlib import Path
+            local_rank = int(os.getenv("LOCAL_RANK", "0"))
+            if local_rank == 0:
+                # download only on first process
+                cached_repo_dir = snapshot_download(
+                    model_name,
+                    local_files_only=is_offline_mode(),
+                    cache_dir=os.getenv("HF_HOME", None),
+                    ignore_patterns=["*.safetensors", "*.msgpack", "*.h5"],
+                )
+                file_list = [str(entry) for entry in Path(cached_repo_dir).rglob("*.[bp][it][n]") if entry.is_file()]
+                data = {"type": "ds_model", "checkpoints": file_list, "version": 1.0}
+                os.makedirs(os.path.join(class_tmpdir, model_name), exist_ok=True)
+                json.dump(data, open(os.path.join(class_tmpdir, model_name, "ds_inference_config.json"), "w"))
+            dist.barrier()
+
+        world_size = int(os.getenv("WORLD_SIZE", "1"))
+        inf_config = {
+            "replace_with_kernel_inject": False,
+            "tensor_parallel": {
+                "tp_size": world_size
+            },
+            "checkpoint": os.path.join(class_tmpdir, model_name, "ds_inference_config.json"),
+        }
+
+        write_checkpoints_json(model_name, class_tmpdir)
+
+        # Load model on meta tensors
+        model_config = AutoConfig.from_pretrained(model_name)
+        # Note that we use half precision to load initially, even for int8
+        with deepspeed.OnDevice(dtype=torch.bfloat16, device="meta"):
+            model = AutoModelForCausalLM.from_config(model_config, torch_dtype=torch.bfloat16)
+        model = model.eval()
+        model = deepspeed.init_inference(model, config=inf_config)

tests/unit/inference/test_inference.py

@@ -0,0 +1,743 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+
+import os
+import time
+
+import deepspeed
+import torch
+
+from packaging import version as pkg_version
+from torch import nn
+from transformers import pipeline
+from transformers.models.t5.modeling_t5 import T5Block
+from transformers.models.roberta.modeling_roberta import RobertaLayer
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.git_version_info import torch_info
+from deepspeed.model_implementations import DeepSpeedTransformerInference
+from deepspeed.ops.op_builder import InferenceBuilder
+from deepspeed.ops.op_builder import OpBuilder
+
+from unit.common import DistributedTest
+
+rocm_version = OpBuilder.installed_rocm_version()
+if rocm_version != (0, 0):
+    pytest.skip("skip inference tests on rocm for now", allow_module_level=True)
+
+_bert_models = [
+    "google-bert/bert-base-cased",
+    "google-bert/bert-base-uncased",
+    "google-bert/bert-large-cased",
+    "google-bert/bert-large-uncased",
+    "google-bert/bert-base-multilingual-cased",
+    "google-bert/bert-base-multilingual-uncased",
+    "deepset/minilm-uncased-squad2",
+    "cross-encoder/ms-marco-MiniLM-L-12-v2",
+    "dslim/bert-base-NER",
+    "google-bert/bert-large-uncased-whole-word-masking-finetuned-squad",
+    "distilbert/distilbert-base-cased-distilled-squad",
+]
+_roberta_models = [
+    "FacebookAI/roberta-large",
+    "FacebookAI/roberta-base",
+    "deepset/roberta-base-squad2",
+    "j-hartmann/emotion-english-distilroberta-base",
+    "Jean-Baptiste/roberta-large-ner-english",
+]
+_gpt_models = [
+    "openai-community/gpt2",
+    "distilbert/distilgpt2",
+    "Norod78/hebrew-bad_wiki-gpt_neo-tiny",
+    "EleutherAI/gpt-j-6b",
+    "EleutherAI/pythia-70m-deduped",
+    "bigscience/bloom-560m",
+]
+_opt_models = [
+    "facebook/opt-125m",  # 125m, 1.7B, ..., 175B variants have the same model architecture.
+    "facebook/opt-350m",  # 350m applies layer norm after attention layer which is different than other variants.
+]
+_test_models = set(_bert_models + _roberta_models + _gpt_models + _opt_models)
+_test_tasks = [
+    "fill-mask", "question-answering", "text-classification", "token-classification", "text-generation",
+    "text2text-generation", "summarization", "translation"
+]
+_test_model_list = _bert_models + _roberta_models + _gpt_models + _opt_models
+
+_model_pipeline_tags = {
+    "google-bert/bert-base-cased": "fill-mask",
+    "google-bert/bert-base-uncased": "fill-mask",
+    "google-bert/bert-large-cased": "fill-mask",
+    "google-bert/bert-large-uncased": "fill-mask",
+    "google-bert/bert-base-multilingual-cased": "fill-mask",
+    "google-bert/bert-base-multilingual-uncased": "fill-mask",
+    "deepset/minilm-uncased-squad2": "question-answering",
+    "cross-encoder/ms-marco-MiniLM-L-12-v2": "text-ranking",
+    "dslim/bert-base-NER": "token-classification",
+    "google-bert/bert-large-uncased-whole-word-masking-finetuned-squad": "question-answering",
+    "distilbert/distilbert-base-cased-distilled-squad": "question-answering",
+    "FacebookAI/roberta-large": "fill-mask",
+    "FacebookAI/roberta-base": "fill-mask",
+    "deepset/roberta-base-squad2": "question-answering",
+    "j-hartmann/emotion-english-distilroberta-base": "text-classification",
+    "Jean-Baptiste/roberta-large-ner-english": "token-classification",
+    "openai-community/gpt2": "text-generation",
+    "distilbert/distilgpt2": "text-generation",
+    "Norod78/hebrew-bad_wiki-gpt_neo-tiny": "text-generation",
+    "EleutherAI/gpt-j-6b": "text-generation",
+    "EleutherAI/pythia-70m-deduped": "text-generation",
+    "bigscience/bloom-560m": "text-generation",
+    "facebook/opt-125m": "text-generation",
+    "facebook/opt-350m": "text-generation",
+}
+
+# Build the deterministic model/task matrix without querying Hugging Face during collection.
+_model_w_tasks = [(model, task) for model in _test_model_list for task in [_model_pipeline_tags[model]]
+                  if task in _test_tasks]
+
+# Assign to pytest variables for testing
+pytest.model_w_tasks = _model_w_tasks
+pytest.mt_names = [f"{m}-{t}" for m, t in pytest.model_w_tasks]
+
+
+@pytest.fixture(scope="module", autouse=True)
+def verify_models():
+    models_without_tags = _test_models.difference(_model_pipeline_tags)
+    if models_without_tags:
+        pytest.fail(f"Model(s) missing a test pipeline tag: {models_without_tags}")
+
+    unknown_models = set(_model_pipeline_tags).difference(_test_models)
+    if unknown_models:
+        pytest.fail(f"Model pipeline tag mapping contains unknown model(s): {unknown_models}")
+
+    missing_tags = {model for model, tag in _model_pipeline_tags.items() if not tag}
+    if missing_tags:
+        pytest.fail(f"Model pipeline tag mapping contains empty tag(s): {missing_tags}")
+
+
+""" Fixtures for inference config """
+
+
+@pytest.fixture(params=pytest.model_w_tasks, ids=pytest.mt_names)
+def model_w_task(request):
+    return request.param
+
+
+@pytest.fixture(params=[torch.float, torch.half], ids=["fp32", "fp16"])
+def dtype(request):
+    return request.param
+
+
+@pytest.fixture(params=[True, False], ids=["CG", "noCG"])
+def enable_cuda_graph(request):
+    return request.param
+
+
+@pytest.fixture(params=[True, False], ids=["Triton", "noTriton"])
+def enable_triton(request):
+    return request.param
+
+
+@pytest.fixture(params=[1, 2], ids=["ws1", "ws2"])
+def world_size(request):
+    return request.param
+
+
+""" Fixtures for running query """
+
+
+@pytest.fixture
+def query(model_w_task):
+    model, task = model_w_task
+    angle_bracket_mask_models = ["roberta", "camembert", "esm", "ibert", "luke", "mpnet", "yoso", "mpnet"]
+
+    if task == "fill-mask":
+        if any(map(lambda x: x in model, angle_bracket_mask_models)):
+            return "Hello I'm a <mask> model."
+        else:
+            return "Hell I'm a [MASK] model."
+    elif task == "question-answering":
+        return {
+            "question": "What's my name?",
+            "context": "My name is Clara and I live in Berkeley",
+        }
+    elif task == "text-classification":
+        return "DeepSpeed is the greatest"
+    elif task == "token-classification":
+        return "My name is jean-baptiste and I live in montreal."
+    elif task == "text-generation":
+        return "DeepSpeed is the greatest"
+    elif task == "text2text-generation":
+        return "Is this review positive or negative? Review: this is the best cast iron skillet you will ever buy"
+    elif task == "translation" or task == "summarization":
+        return "Hello, my dog is cute"
+    else:
+        NotImplementedError(f'query for task "{task}" is not implemented')
+
+
+@pytest.fixture
+def inf_kwargs(model_w_task):
+    model, task = model_w_task
+    if task == "text-generation":
+        if model == "EleutherAI/gpt-j-6b":
+            # This model on V100 is hitting memory problems that limit the number of output tokens
+            return {"do_sample": False, "temperature": 1.0, "max_length": 12}
+        return {"do_sample": False, "temperature": 1.0, "max_length": 20}
+    else:
+        return {}
+
+
+""" Assertion fixture for verifying model outputs """
+
+
+def fill_mask_assert(x, y):
+    return set(res["token_str"] for res in x) == set(res["token_str"] for res in y)
+
+
+def question_answering_assert(x, y):
+    return x["answer"] == y["answer"]
+
+
+def text_classification_assert(x, y):
+    return set(res["label"] for res in x) == set(res["label"] for res in y)
+
+
+def token_classification_assert(x, y):
+    return set(ent["word"] for ent in x) == set(ent["word"] for ent in y)
+
+
+def text_generation_assert(x, y):
+    return set(res["generated_text"] for res in x) == set(res["generated_text"] for res in y)
+
+
+def text2text_generation_assert(x, y):
+    return set(res["generated_text"] for res in x) == set(res["generated_text"] for res in y)
+
+
+def translation_assert(x, y):
+    return set(res["translation_text"] for res in x) == set(res["translation_text"] for res in y)
+
+
+def summarization_assert(x, y):
+    return set(res["summary_text"] for res in x) == set(res["summary_text"] for res in y)
+
+
+@pytest.fixture
+def assert_fn(model_w_task):
+    model, task = model_w_task
+    assert_fn_dict = {
+        "fill-mask": fill_mask_assert,
+        "question-answering": question_answering_assert,
+        "text-classification": text_classification_assert,
+        "token-classification": token_classification_assert,
+        "text-generation": text_generation_assert,
+        "text2text-generation": text2text_generation_assert,
+        "translation": translation_assert,
+        "summarization": summarization_assert
+    }
+    assert_fn = assert_fn_dict.get(task, None)
+    if assert_fn is None:
+        NotImplementedError(f'assert_fn for task "{task}" is not implemented')
+    return assert_fn
+
+
+# Used to verify DeepSpeed kernel injection worked with a model
+def check_injection(model):
+
+    def verify_injection(module):
+        for child in module.children():
+            if isinstance(child, nn.ModuleList):
+                assert isinstance(child[0], DeepSpeedTransformerInference),\
+                    "DeepSpeed-Inference Transformer kernels has not been injected in the model"
+                break
+            else:
+                verify_injection(child)
+
+    verify_injection(model)
+
+
+# Used to Get Device name
+def getDeviceId(local_rank):
+    device = torch.device(f"{get_accelerator().device_name(local_rank)}")
+    return device
+
+
+# Verify that test is valid
+def validate_test(model_w_task, dtype, enable_cuda_graph, enable_triton):
+    model, task = model_w_task
+    msg = ""
+    if enable_cuda_graph and (torch_info["cuda_version"] == "0.0"):
+        msg = "CUDA not detected, cannot use CUDA Graph"
+    elif enable_cuda_graph and pkg_version.parse(torch.__version__) < pkg_version.parse("1.10"):
+        msg = "CUDA Graph is only available in torch versions >= 1.10"
+    elif "gpt-j-6b" in model:
+        if dtype != torch.half:
+            msg = f"Not enough GPU memory to run {model} with dtype {dtype}"
+        elif enable_cuda_graph:
+            msg = f"Not enough GPU memory to run {model} with CUDA Graph enabled"
+    elif "gpt-neox-20b" in model:  # TODO: remove this when neox issues resolved
+        msg = "Skipping gpt-neox-20b for now"
+    elif ("gpt-neox-20b" in model) and (dtype != torch.half):
+        msg = f"Not enough GPU memory to run {model} with dtype {dtype}"
+    elif ("bloom" in model) and (dtype != torch.half):
+        msg = f"Bloom models only support half precision, cannot use dtype {dtype}"
+    elif (model not in _bert_models + _roberta_models) and enable_cuda_graph:
+        msg = "Non bert/roberta models do no support CUDA Graph"
+    elif enable_triton and not (dtype in [torch.half]):
+        msg = "Triton is for fp16"
+    elif enable_triton and not deepspeed.HAS_TRITON:
+        msg = "triton needs to be installed for the test"
+    elif (model not in _bert_models + _roberta_models) and enable_triton:
+        msg = "Triton kernels do not support Non bert/roberta models yet"
+
+    # These should be removed once we fix several inference tests failing
+    if model in [
+            "EleutherAI/pythia-70m-deduped", "distilbert/distilbert-base-cased-distilled-squad", "EleutherAI/gpt-j-6b"
+    ]:
+        msg = "Test is currently broken"
+    return msg
+
+
+@pytest.mark.inference
+class TestModelTask(DistributedTest):
+    world_size = 1
+
+    def test(
+        self,
+        model_w_task,
+        dtype,
+        enable_cuda_graph,
+        enable_triton,
+        query,
+        inf_kwargs,
+        assert_fn,
+        perf_meas=True,
+    ):
+        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph, enable_triton)
+        if invalid_test_msg:
+            pytest.skip(invalid_test_msg)
+
+        if dtype not in get_accelerator().supported_dtypes():
+            pytest.skip(f"Acceleraor {get_accelerator().device_name()} does not support {dtype}.")
+
+        if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
+            pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+        model, task = model_w_task
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+
+        # Load the model on CPU first to avoid OOM for large models @fp32
+        pipe = pipeline(task, model=model, device=torch.device("cpu"), framework="pt")
+        if dtype == torch.half:
+            pipe.model.half()
+
+        # Switch device to GPU after converting to half
+        device = torch.device(get_accelerator().device_name(local_rank))
+        pipe.device = device
+        pipe.model.to(device)
+
+        # Warm-up queries for perf measurement
+        #for i in range(10):
+        #    _ = pipe(query, **inf_kwargs)
+        get_accelerator().synchronize()
+        start = time.time()
+        bs_output = pipe(query, **inf_kwargs)
+        get_accelerator().synchronize()
+        bs_time = time.time() - start
+
+        args = {
+            'mp_size': 1,
+            'dtype': dtype,
+            'replace_with_kernel_inject': True,
+            'enable_cuda_graph': enable_cuda_graph,
+            'use_triton': enable_triton,
+            'triton_autotune': False,
+        }
+        if pipe.tokenizer.model_max_length < deepspeed.ops.transformer.inference.config.DeepSpeedInferenceConfig(
+        ).max_out_tokens:
+            args.update({'max_out_tokens': pipe.tokenizer.model_max_length})
+        pipe.model = deepspeed.init_inference(pipe.model, **args)
+        check_injection(pipe.model)
+        # Warm-up queries for perf measurement
+        #for i in range(10):
+        #    _ = pipe(query, **inf_kwargs)
+        get_accelerator().synchronize()
+        start = time.time()
+        ds_output = pipe(query, **inf_kwargs)
+        get_accelerator().synchronize()
+        ds_time = time.time() - start
+
+        if perf_meas:
+            print(
+                f"model={model}, task={task}, dtype={dtype}, cuda_graph={enable_cuda_graph}, triton={enable_triton}, bs_time={bs_time}, ds_time={ds_time}"
+            )
+
+        # facebook/opt* and some bigscient/bloom* models are not matching
+        # baseline exactly, adding an exception to them for now
+        if ("opt" in model) or ("bloom" in model):
+            bs_output = pipe(query, **inf_kwargs)
+
+        # These performance tests are only measuring the time for a single
+        # inference request, we just want to check that performance isn't terrible
+        #assert ds_time <= (bs_time * 1.1)
+
+        assert assert_fn(bs_output, ds_output)
+
+
+@pytest.mark.seq_inference
+@pytest.mark.parametrize("model_w_task", [("EleutherAI/gpt-neo-1.3B", "text-generation"),
+                                          ("EleutherAI/gpt-neox-20b", "text-generation"),
+                                          ("bigscience/bloom-3b", "text-generation"),
+                                          ("EleutherAI/gpt-j-6b", "text-generation")],
+                         ids=["gpt-neo", "gpt-neox", "bloom", "gpt-j"])
+class TestMPSize(DistributedTest):
+    world_size = 2
+
+    def test(
+        self,
+        model_w_task,
+        dtype,
+        query,
+        inf_kwargs,
+        assert_fn,
+    ):
+        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
+        if invalid_test_msg:
+            pytest.skip(invalid_test_msg)
+
+        if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
+            pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+        model, task = model_w_task
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+
+        # We have to load these large models on CPU with pipeline because not
+        # enough GPU memory
+        pipe = pipeline(task, model=model, device=torch.device("cpu"), framework="pt")
+        bs_output = pipe(query, **inf_kwargs)
+
+        pipe.model = deepspeed.init_inference(pipe.model,
+                                              mp_size=self.world_size,
+                                              dtype=dtype,
+                                              replace_with_kernel_inject=True)
+        check_injection(pipe.model)
+        # Switch device to GPU so that input tensors are not on CPU
+        pipe.device = torch.device(get_accelerator().device_name(local_rank))
+        ds_output = pipe(query, **inf_kwargs)
+
+        print(local_rank, "baseline", bs_output)
+        print(local_rank, "deepspeed", ds_output)
+        assert assert_fn(bs_output, ds_output)
+
+
+@pytest.mark.inference
+@pytest.mark.parametrize("model_w_task", [("openai-community/gpt2", "text-generation")], ids=["gpt2"])
+class TestLowCpuMemUsage(DistributedTest):
+    world_size = 1
+
+    def test(
+        self,
+        model_w_task,
+        query,
+        inf_kwargs,
+        assert_fn,
+    ):
+        model, task = model_w_task
+        dtype = torch.float16
+        if dtype not in get_accelerator().supported_dtypes():
+            pytest.skip(f"Acceleraor {get_accelerator().device_name()} does not support {dtype}.")
+
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        device = getDeviceId(local_rank)
+        pipe = pipeline(task, model=model, model_kwargs={"low_cpu_mem_usage": True}, device=device, framework="pt")
+        bs_output = pipe(query, **inf_kwargs)
+        pipe.model = deepspeed.init_inference(pipe.model,
+                                              mp_size=self.world_size,
+                                              dtype=dtype,
+                                              replace_method="auto",
+                                              replace_with_kernel_inject=True)
+
+        ds_output = pipe(query, **inf_kwargs)
+
+        assert assert_fn(bs_output, ds_output)
+
+
+@pytest.mark.seq_inference
+@pytest.mark.parametrize(
+    "model_w_task, injection_policy",
+    [
+        (("google/t5-v1_1-small", "text2text-generation"), {
+            T5Block: ('SelfAttention.o', 'EncDecAttention.o', 'DenseReluDense.wo')
+        }),
+        (("FacebookAI/roberta-large", "fill-mask"), {
+            RobertaLayer: ('output.dense')
+        }),
+    ],
+    ids=["t5", "roberta"],
+)
+@pytest.mark.parametrize("dtype", [torch.float], ids=["fp32"])
+class TestInjectionPolicy(DistributedTest):
+
+    def test(self, model_w_task, injection_policy, query, inf_kwargs, assert_fn, dtype, world_size):
+        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
+        if invalid_test_msg:
+            pytest.skip(invalid_test_msg)
+
+        model, task = model_w_task
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+
+        pipe = pipeline(task,
+                        model=model,
+                        device=torch.device(get_accelerator().device_name(local_rank)),
+                        framework="pt")
+        bs_output = pipe(query, **inf_kwargs)
+
+        pipe.model = deepspeed.init_inference(pipe.model,
+                                              mp_size=world_size,
+                                              dtype=dtype,
+                                              injection_policy=injection_policy)
+        ds_output = pipe(query, **inf_kwargs)
+
+        print(local_rank, "baseline", bs_output)
+        print(local_rank, "deepspeed", ds_output)
+        assert assert_fn(bs_output, ds_output)
+
+
+@pytest.mark.seq_inference
+@pytest.mark.parametrize("model_w_task", [("Felladrin/Llama-160M-Chat-v1", "text-generation")], ids=["llama"])
+@pytest.mark.parametrize("dtype", [torch.half], ids=["fp16"])
+class TestLlamaInjection(DistributedTest):
+    world_size = 1
+
+    def test(self, model_w_task, dtype, query, inf_kwargs, assert_fn):
+        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
+        if invalid_test_msg:
+            pytest.skip(invalid_test_msg)
+
+        if dtype not in get_accelerator().supported_dtypes():
+            pytest.skip(f"Accelerator {get_accelerator().device_name()} does not support {dtype}.")
+
+        if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
+            pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+        model, task = model_w_task
+
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        device = torch.device(get_accelerator().device_name(local_rank))
+
+        pipe = pipeline(task,
+                        model=model,
+                        device=torch.device("cpu"),
+                        model_kwargs={"low_cpu_mem_usage": True},
+                        framework="pt")
+
+        if dtype == torch.half:
+            pipe.model.half()
+
+        pipe.device = device
+        pipe.model.to(device)
+        bs_output = pipe(query, **inf_kwargs)
+
+        try:
+            pipe.model = deepspeed.init_inference(pipe.model,
+                                                  mp_size=self.world_size,
+                                                  dtype=dtype,
+                                                  replace_with_kernel_inject=True)
+            check_injection(pipe.model)
+        except AttributeError as e:
+            if "'LlamaAttention' object has no attribute 'num_heads'" in str(e):
+                pytest.skip("Skipping due to transformers version compatibility issue with self-attention")
+            raise e
+
+        ds_output = pipe(query, **inf_kwargs)
+
+        print(local_rank, "baseline", bs_output)
+        print(local_rank, "deepspeed", ds_output)
+        # Llama models are not matching baseline exactly
+        # We skip the result check for now, since this is irrelevant to this test
+        # assert assert_fn(bs_output, ds_output)
+
+
+@pytest.mark.seq_inference
+@pytest.mark.parametrize('keep_module_on_host', [True, False])
+@pytest.mark.parametrize(
+    "model_w_task",
+    [("Helsinki-NLP/opus-mt-en-de", "translation"), ("Salesforce/codegen-350M-mono", "text-generation")],
+    ids=["marian", "codegen"],  #codegen has fusedqkv weight.
+)
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"])
+class TestAutoTensorParallelism(DistributedTest):
+    world_size = [2]
+
+    def test(
+        self,
+        model_w_task,
+        query,
+        inf_kwargs,
+        assert_fn,
+        dtype,
+        keep_module_on_host,
+    ):
+        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
+        if invalid_test_msg:
+            pytest.skip(invalid_test_msg)
+
+        model, task = model_w_task
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        world_size = int(os.getenv("WORLD_SIZE", "2"))
+
+        if dtype not in get_accelerator().supported_dtypes():
+            pytest.skip(f"Acceleraor {get_accelerator().device_name()} does not support {dtype}.")
+
+        if model == "Salesforce/codegen-350M-mono":
+            pytest.skip("Disable Codegen model due to slight result difference")
+            #TODO: re-enable this test once we have a fix for the slight result difference
+
+        pipe = pipeline(task,
+                        model=model,
+                        device=torch.device(get_accelerator().device_name(local_rank)),
+                        framework="pt")
+        bs_output = pipe(query, **inf_kwargs)
+
+        pipe.model = deepspeed.init_inference(pipe.model,
+                                              mp_size=world_size,
+                                              dtype=dtype,
+                                              keep_module_on_host=keep_module_on_host)
+        ds_output = pipe(query, **inf_kwargs)
+
+        print(local_rank, "baseline", bs_output)
+        print(local_rank, "deepspeed", ds_output)
+        assert assert_fn(bs_output, ds_output)
+
+        if keep_module_on_host:
+            for name, param in model.named_parameters():
+                assert param.device == torch.device('cpu'), f"keep_module_on_host is on but param {name} is not on cpu"
+
+    @pytest.mark.world_size(3)
+    def test_odd_world_size(
+        self,
+        model_w_task,
+        query,
+        inf_kwargs,
+        assert_fn,
+        dtype,
+        keep_module_on_host,
+    ):
+        invalid_test_msg = validate_test(model_w_task, dtype, enable_cuda_graph=False, enable_triton=False)
+        if invalid_test_msg:
+            pytest.skip(invalid_test_msg)
+
+        model, task = model_w_task
+        if model == "Salesforce/codegen-350M-mono":
+            pytest.skip("codegen does not supported by odd world_size")
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        world_size = int(os.getenv("WORLD_SIZE", "3"))
+
+        pipe = pipeline(task,
+                        model=model,
+                        device=torch.device(get_accelerator().device_name(local_rank)),
+                        framework="pt")
+        bs_output = pipe(query, **inf_kwargs)
+
+        pipe.model = deepspeed.init_inference(pipe.model,
+                                              mp_size=world_size,
+                                              dtype=dtype,
+                                              keep_module_on_host=keep_module_on_host)
+        ds_output = pipe(query, **inf_kwargs)
+
+        print(local_rank, "baseline", bs_output)
+        print(local_rank, "deepspeed", ds_output)
+        assert assert_fn(bs_output, ds_output)
+
+        if keep_module_on_host:
+            for name, param in model.named_parameters():
+                assert param.device == torch.device('cpu'), f"keep_module_on_host is on but param {name} is not on cpu"
+
+
+@pytest.mark.nightly
+@pytest.mark.parametrize(
+    "model_family, model_name",
+    (
+        ["gpt2", "EleutherAI/gpt-neo-2.7B"],
+        #["gpt2", "EleutherAI/gpt-j-6b"], # Causing OOM for this test
+        ["gpt2", "openai-community/gpt2-xl"],
+    ),
+)
+@pytest.mark.parametrize("task", ["lambada_standard"])
+class TestLMCorrectness(DistributedTest):
+    world_size = 1
+    exec_timeout = 1200  # Give these tests longer to complete
+
+    def test(self, model_family, model_name, task):
+        # imports here to avoid import errors when pytest collects tests
+        import lm_eval
+        import lm_eval.models
+        import lm_eval.tasks
+        import lm_eval.evaluator
+
+        # The bootstrap_stderr function in lm_eval.metrics uses a
+        # multiprocessing Pool to increase performance. Since we use a Pool for
+        # our distributed tests and cannot nest Pools, we must redefine and
+        # patch this function with a version that does not use Pool.
+        def no_pool_bootstrap_stderr(f, xs, iters):
+            from lm_eval.metrics import _bootstrap_internal
+            from lm_eval.metrics import sample_stddev
+            res = []
+            chunk_size = min(1000, iters)
+            for i in range(iters // chunk_size):
+                res.extend(_bootstrap_internal(f, chunk_size)((i, xs)))
+            return sample_stddev(res)
+
+        lm_eval.metrics.bootstrap_stderr = no_pool_bootstrap_stderr
+
+        local_rank = os.getenv("LOCAL_RANK", "0")
+        device = torch.device(get_accelerator().device_name(local_rank))
+        dtype = torch.float
+        task_dict = lm_eval.tasks.get_task_dict([task])
+
+        if 'gpt-j-6b' in model_name:
+            dtype = torch.half
+            lm = lm_eval.models.get_model(model_family).create_from_arg_string(f"pretrained={model_name}",
+                                                                               {"device": "cpu"})
+            setattr(lm, model_family, getattr(lm, model_family).half().to(device))
+            lm._device = device
+        else:
+            if get_accelerator().device_name() == 'hpu':
+                #lm_eval not supporting HPU device, so get model with CPU and move it to HPU.
+                lm = lm_eval.models.get_model(model_family).create_from_arg_string(f"pretrained={model_name}",
+                                                                                   {"device": "cpu"})
+                setattr(lm, model_family, getattr(lm, model_family).to(device))
+                lm._device = device
+            else:
+                lm = lm_eval.models.get_model(model_family).create_from_arg_string(
+                    f"pretrained={model_name}", {"device": get_accelerator().device_name()})
+
+        get_accelerator().synchronize()
+        start = time.time()
+        bs_output = lm_eval.evaluator.evaluate(lm=lm, task_dict=task_dict)
+        get_accelerator().synchronize()
+        bs_time = time.time() - start
+
+        getattr(lm, model_family).to("cpu")
+        ds_model = deepspeed.init_inference(
+            getattr(lm, model_family),
+            mp_size=1,
+            dtype=dtype,
+            replace_with_kernel_inject=True,
+            enable_cuda_graph=False,
+        )
+        check_injection(ds_model)
+        setattr(lm, model_family, ds_model)
+        get_accelerator().synchronize()
+        start = time.time()
+        ds_output = lm_eval.evaluator.evaluate(lm=lm, task_dict=task_dict)
+        get_accelerator().synchronize()
+        ds_time = time.time() - start
+
+        ppl_diff = abs(bs_output["results"][task]["ppl"] - ds_output["results"][task]["ppl"])
+        #assert ds_time <= bs_time
+        assert ppl_diff < 0.01

tests/unit/inference/test_inference_config.py

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+import deepspeed
+from unit.common import DistributedTest
+from unit.simple_model import create_config_from_dict
+
+
+@pytest.mark.inference
+class TestInferenceConfig(DistributedTest):
+    world_size = 1
+
+    def test_overlap_kwargs(self):
+        config = {"replace_with_kernel_inject": True, "dtype": torch.float32}
+        kwargs = {"replace_with_kernel_inject": True}
+
+        engine = deepspeed.init_inference(torch.nn.Module(), config=config, **kwargs)
+        assert engine._config.replace_with_kernel_inject
+
+    def test_overlap_kwargs_conflict(self):
+        config = {"replace_with_kernel_inject": True}
+        kwargs = {"replace_with_kernel_inject": False}
+
+        with pytest.raises(ValueError):
+            engine = deepspeed.init_inference(torch.nn.Module(), config=config, **kwargs)
+
+    def test_kwargs_and_config(self):
+        config = {"replace_with_kernel_inject": True}
+        kwargs = {"dtype": torch.float32}
+
+        engine = deepspeed.init_inference(torch.nn.Module(), config=config, **kwargs)
+        assert engine._config.replace_with_kernel_inject
+        assert engine._config.dtype == kwargs["dtype"]
+
+    def test_json_config(self, tmpdir):
+        config = {"replace_with_kernel_inject": True, "dtype": "torch.float32"}
+        config_json = create_config_from_dict(tmpdir, config)
+
+        engine = deepspeed.init_inference(torch.nn.Module(), config=config_json)
+        assert engine._config.replace_with_kernel_inject

tests/unit/inference/test_model_profiling.py

@@ -0,0 +1,61 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import time
+import pytest
+import torch
+import deepspeed
+from transformers import pipeline
+from unit.common import DistributedTest
+from deepspeed.accelerator import get_accelerator
+from deepspeed.ops.op_builder import InferenceBuilder
+
+if not deepspeed.ops.__compatible_ops__[InferenceBuilder.NAME]:
+    pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
+
+if torch.half not in get_accelerator().supported_dtypes():
+    pytest.skip(f"fp16 not supported, valid dtype: {get_accelerator().supported_dtypes()}", allow_module_level=True)
+
+
+@pytest.mark.inference
+@pytest.mark.parametrize("use_cuda_events", [True, False])
+@pytest.mark.parametrize("enable_cuda_graph", [True, False])
+class TestModelProfiling(DistributedTest):
+    world_size = 1
+
+    def test(self, enable_cuda_graph, use_cuda_events):
+        task = "fill-mask"
+        model = "bert-base-cased"
+        dtype = torch.float16
+        query = "I am a [MASK] model"
+
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        world_size = int(os.getenv("WORLD_SIZE", "1"))
+
+        pipe = pipeline(task, model, framework="pt", device=get_accelerator().device_name(local_rank))
+        pipe.model = deepspeed.init_inference(pipe.model,
+                                              dtype=dtype,
+                                              mp_size=world_size,
+                                              replace_with_kernel_inject=True,
+                                              enable_cuda_graph=enable_cuda_graph)
+        pipe.model.profile_model_time(use_cuda_events=use_cuda_events)
+
+        e2e_times = []
+        model_times = []
+        for _ in range(10):
+            get_accelerator().synchronize()
+            start = time.perf_counter_ns()
+
+            r = pipe(query)
+
+            get_accelerator().synchronize()
+            end = time.perf_counter_ns()
+
+            e2e_times.append((end - start) / 1e6)  # convert ns to ms
+            model_times.extend(pipe.model.model_times())
+
+        for e2e_t, model_t in zip(e2e_times, model_times):
+            assert e2e_t >= model_t

tests/unit/inference/test_stable_diffusion.py

@@ -0,0 +1,48 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import os
+import torch
+import pytest
+import deepspeed
+import numpy
+from unit.common import DistributedTest
+from deepspeed.accelerator import get_accelerator
+
+
+# Setup for these models is different from other pipelines, so we add a separate test
+@pytest.mark.stable_diffusion
+class TestStableDiffusion(DistributedTest):
+    world_size = 1
+
+    def test(self):
+        from diffusers import DiffusionPipeline
+        from image_similarity_measures.quality_metrics import rmse
+        dev = get_accelerator().device_name()
+        generator = torch.Generator(device=dev)
+        seed = 0xABEDABE7
+        generator.manual_seed(seed)
+        prompt = "a dog on a rocket"
+        model = "prompthero/midjourney-v4-diffusion"
+        local_rank = int(os.getenv("LOCAL_RANK", "0"))
+        device = torch.device(f"{dev}:{local_rank}")
+        pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.half)
+        pipe = pipe.to(device)
+        baseline_image = pipe(prompt, guidance_scale=7.5, generator=generator).images[0]
+
+        pipe = deepspeed.init_inference(
+            pipe,
+            mp_size=1,
+            dtype=torch.half,
+            replace_with_kernel_inject=True,
+            enable_cuda_graph=True,
+        )
+        generator.manual_seed(seed)
+        deepspeed_image = pipe(prompt, guidance_scale=7.5, generator=generator).images[0]
+
+        rmse_value = rmse(org_img=numpy.asarray(baseline_image), pred_img=numpy.asarray(deepspeed_image))
+
+        # RMSE threshold value is arbitrary, may need to adjust as needed
+        assert rmse_value <= 0.01

tests/unit/inference/v2/__init__.py

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

tests/unit/inference/v2/inference_test_utils.py

@@ -0,0 +1,46 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Tuple
+
+import torch
+from deepspeed.accelerator import get_accelerator
+
+TOLERANCES = None
+
+
+def get_tolerances():
+    global TOLERANCES
+    if TOLERANCES is None:
+        TOLERANCES = {torch.float32: (5e-4, 5e-5), torch.float16: (3e-2, 2e-3)}
+        if get_accelerator().is_bf16_supported():
+            # Note: BF16 tolerance is higher than FP16 because of the lower precision (7 (+1) bits vs
+            # 10 (+1) bits)
+            TOLERANCES[torch.bfloat16] = (4.8e-1, 3.2e-2)
+    return TOLERANCES
+
+
+DTYPES = None
+
+
+def get_dtypes(include_float=True):
+    global DTYPES
+    if DTYPES is None:
+        DTYPES = [torch.float16, torch.float32] if include_float else [torch.float16]
+        try:
+            if get_accelerator().is_bf16_supported():
+                DTYPES.append(torch.bfloat16)
+        except (AssertionError, AttributeError):
+            pass
+    return DTYPES
+
+
+def allclose(x, y, tolerances: Tuple[int, int] = None):
+    assert x.dtype == y.dtype
+    if tolerances is None:
+        rtol, atol = get_tolerances()[x.dtype]
+    else:
+        rtol, atol = tolerances
+    return torch.allclose(x, y, rtol=rtol, atol=atol)

tests/unit/inference/v2/kernels/__init__.py

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

tests/unit/inference/v2/kernels/core_ops/__init__.py

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

tests/unit/inference/v2/kernels/core_ops/test_bias_activation.py

@@ -0,0 +1,101 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum
+from deepspeed.inference.v2.kernels.core_ops import CUDABiasActivation
+from ....v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_bias_act_implementation(input: torch.Tensor, bias: Optional[torch.Tensor],
+                                      act_type: ActivationType) -> torch.Tensor:
+    bias_func_map = {
+        ActivationType.RELU: torch.nn.functional.relu,
+        ActivationType.GELU: torch.nn.functional.gelu,
+        ActivationType.SILU: torch.nn.functional.silu,
+        ActivationType.IDENTITY: lambda x: x,
+    }
+
+    dtype = input.dtype
+    input_f = input.to(torch.float32)
+    if bias is not None:
+        bias_f = bias.to(torch.float32)
+        output_f = input_f + bias_f
+    else:
+        output_f = input_f
+    output_f = bias_func_map[act_type](output_f)
+
+    return output_f.to(dtype)
+
+
+def _bias_activation_test_helper(tokens: int,
+                                 channels: int,
+                                 act_fn: ActivationType,
+                                 dtype: DtypeEnum,
+                                 use_bias: bool = True) -> None:
+    """
+    Fully parameterized testing entry point.
+    """
+    # Input vals
+    input_tensor = torch.randn((tokens, channels), dtype=dtype.value, device=get_accelerator().current_device_name())
+    if use_bias:
+        bias = torch.randn((channels), dtype=dtype.value, device=get_accelerator().current_device_name())
+    else:
+        bias = None
+
+    # Reference output
+    ref_output = reference_bias_act_implementation(input_tensor, bias, act_fn)
+
+    bias_act = CUDABiasActivation(channels, dtype, act_fn)
+
+    # New output
+    ds_tensor = input_tensor.clone()
+    bias_act(ds_tensor, bias)
+
+    # Check
+    assert allclose(ds_tensor, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, channels", [(1, 4096), (37, 2048), (112, 14432), (1024, 6144)])
+@pytest.mark.parametrize("dtype", get_dtypes(include_float=False))
+def test_token_channels_permutations(tokens: int, channels: int, dtype: torch.dtype) -> None:
+    """
+    Validate bias activation kernel with different token and channel permutations when using the RELU
+    activation function.
+    """
+    act_fn = ActivationType.RELU
+    dtype = DtypeEnum(dtype)
+    _bias_activation_test_helper(tokens, channels, act_fn, dtype)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("act_fn",
+                         [ActivationType.RELU, ActivationType.GELU, ActivationType.SILU, ActivationType.IDENTITY])
+def test_act_fns(act_fn: ActivationType) -> None:
+    """
+    Validate bias activation kernel with different activation functions.
+    """
+    tokens = 223
+    channels = 4096
+    dtype = DtypeEnum.fp16
+    _bias_activation_test_helper(tokens, channels, act_fn, dtype)
+
+
+@pytest.mark.inference_v2_ops
+def test_no_bias() -> None:
+    """
+    Validate bias activation kernel with no bias.
+    """
+    tokens = 223
+    channels = 4096
+    dtype = DtypeEnum.fp16
+    act_fn = ActivationType.IDENTITY
+    _bias_activation_test_helper(tokens, channels, act_fn, dtype, use_bias=False)

tests/unit/inference/v2/kernels/core_ops/test_blas_linear.py

@@ -0,0 +1,73 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.core_ops import BlasLibLinear
+from ....v2.inference_test_utils import allclose
+
+# Note: only testing with FP16 and BF16 because we use TF32 on Ampere and we don't have a good
+# set of tolerances. Since this is just on top of BLAS though, the test is more about
+# making sure the stride/contiguity is correct and that's data type agnostic.
+
+
+def reference_implementation(hidden_states, weights):
+    return hidden_states @ weights.t()
+
+
+problem_shapes = [
+    (1, 1, 1024, 1024),
+    (1, 1024, 1024, 1024),
+    (2, 1024, 1024, 1024),
+    (1, 128, 768, 3072),
+    (1, 128, 3072, 768),
+    (1, 1024, 8192, 8192),
+    (1, 733, 8192, 32768),
+    (1, 13, 32768, 8192),
+]
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("fp_dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("problem_shape", problem_shapes)
+def test_blas_linear(fp_dtype: torch.dtype, problem_shape: Tuple[int, int, int, int]):
+    batch, seq_len, in_features, out_features = problem_shape
+    hidden_states = torch.randn(batch, seq_len, in_features, dtype=fp_dtype,
+                                device=get_accelerator().current_device()) * 0.1
+    weights = torch.randn(out_features, in_features, dtype=fp_dtype, device=get_accelerator().current_device()) * 0.01
+    ds_output = torch.empty(batch, seq_len, out_features, dtype=fp_dtype, device=get_accelerator().current_device())
+
+    ds_kernel = BlasLibLinear(fp_dtype)
+
+    ds_output = ds_kernel(ds_output, hidden_states, weights)
+    ref_output = reference_implementation(hidden_states, weights)
+
+    assert allclose(ds_output, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("fp_dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("problem_shape", problem_shapes)
+def test_blas_linear_t(fp_dtype: torch.dtype, problem_shape: Tuple[int, int, int, int]):
+    batch, seq_len, in_features, out_features = problem_shape
+    hidden_states = torch.randn(batch, seq_len, in_features, dtype=fp_dtype,
+                                device=get_accelerator().current_device()) * 0.1
+    weights = torch.randn(out_features, in_features, dtype=fp_dtype, device=get_accelerator().current_device()) * 0.01
+    ds_output = torch.empty(batch, seq_len, out_features, dtype=fp_dtype, device=get_accelerator().current_device())
+
+    ds_kernel = BlasLibLinear(fp_dtype)
+
+    # Transpose the weights then revert to the format we expect.
+    weights = weights.t().contiguous()
+    weights = weights.t()
+    ds_output = ds_kernel(ds_output, hidden_states, weights)
+
+    ref_output = reference_implementation(hidden_states, weights)
+
+    assert allclose(ds_output, ref_output)

tests/unit/inference/v2/kernels/core_ops/test_gated_activation.py

@@ -0,0 +1,133 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Iterable, Optional
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.core_ops import CUDAGatedActivation
+from deepspeed.inference.v2.inference_utils import ActivationType
+from ....v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_geglu_implementation(input: torch.Tensor,
+                                   bias: Optional[torch.Tensor] = None,
+                                   act_fn: Optional[ActivationType] = ActivationType.GEGLU) -> torch.Tensor:
+    act_func_map = {
+        ActivationType.ReGLU: torch.nn.functional.relu,
+        ActivationType.GEGLU: lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
+        ActivationType.SiGLU: torch.nn.functional.silu,
+    }
+
+    dtype = input.dtype
+    input = input.to(torch.float32)
+
+    if bias is not None:
+        bias = bias.to(torch.float32)
+        input = input + bias
+
+    act_act = input[..., ::2]
+    act_linear = input[..., 1::2]
+
+    act_act = act_func_map[act_fn](act_act)
+
+    return (act_act * act_linear).to(dtype)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("shape", [(1372, 16384), (2, 743, 22016)])
+@pytest.mark.parametrize("dtype", get_dtypes())
+def test_dtypes(shape: Iterable[int], dtype: torch.dtype) -> None:
+    input_tensor = torch.randn(shape, dtype=dtype, device=get_accelerator().current_device_name())
+
+    # Reference output
+    ref_output = reference_geglu_implementation(input_tensor, act_fn=ActivationType.GEGLU)
+
+    # Build kernel
+    geglu = CUDAGatedActivation(input_tensor.size(-1), input_tensor.dtype, ActivationType.GEGLU)
+
+    # New output
+    output_shape = list(input_tensor.shape)
+    output_shape[-1] //= 2
+    output_tensor = torch.empty(output_shape, dtype=input_tensor.dtype, device=get_accelerator().current_device_name())
+    geglu(output_tensor, input_tensor)
+
+    # Check
+    assert allclose(output_tensor, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("act_fn", [ActivationType.GEGLU, ActivationType.ReGLU, ActivationType.SiGLU])
+def test_act_fn(act_fn: ActivationType) -> None:
+    input_tensor = torch.randn(832, 4096, dtype=torch.float16, device=get_accelerator().current_device())
+
+    # Reference output
+    ref_output = reference_geglu_implementation(input_tensor, act_fn=act_fn)
+
+    cuda_act = CUDAGatedActivation(4096, torch.float16, act_fn)
+
+    # New output
+    output_tensor = torch.empty(832, 2048, dtype=torch.float16, device=get_accelerator().current_device())
+    cuda_act(output_tensor, input_tensor)
+
+    assert allclose(output_tensor, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+def test_act_with_bias():
+    input_tensor = torch.randn(832, 4096, dtype=torch.float16, device=get_accelerator().current_device())
+    bias = torch.randn(4096, dtype=torch.float16, device=get_accelerator().current_device())
+
+    # Reference output
+    ref_output = reference_geglu_implementation(input_tensor, bias=bias, act_fn=ActivationType.GEGLU)
+
+    cuda_act = CUDAGatedActivation(4096, torch.float16, ActivationType.GEGLU)
+
+    # New output
+    output_tensor = torch.empty(832, 2048, dtype=torch.float16, device=get_accelerator().current_device())
+
+    cuda_act(output_tensor, input_tensor, bias)
+
+    assert allclose(output_tensor, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+def test_max_channels():
+    input_tensor = torch.randn(832, 48152, dtype=torch.float16, device=get_accelerator().current_device())
+
+    ref_output = reference_geglu_implementation(input_tensor, act_fn=ActivationType.GEGLU)
+
+    cuda_act = CUDAGatedActivation(48152, torch.float16, ActivationType.GEGLU)
+
+    output_tensor = torch.empty(832, 24076, dtype=torch.float16, device=get_accelerator().current_device())
+    cuda_act(output_tensor, input_tensor)
+
+    assert allclose(output_tensor, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+def test_bad_dtype() -> None:
+    with pytest.raises(ValueError):
+        CUDAGatedActivation(128, torch.int8, ActivationType.GEGLU)
+
+
+@pytest.mark.inference_v2_ops
+def test_bad_act_fn() -> None:
+    with pytest.raises(ValueError):
+        CUDAGatedActivation(128, torch.float16, ActivationType.RELU)
+
+
+@pytest.mark.inference_v2_ops
+def test_bad_alignment() -> None:
+    with pytest.raises(ValueError):
+        CUDAGatedActivation(127, torch.float16, ActivationType.GEGLU)
+
+
+@pytest.mark.inference_v2_ops
+def test_too_many_channels() -> None:
+    with pytest.raises(ValueError):
+        CUDAGatedActivation(49160, torch.float16, ActivationType.GEGLU)

tests/unit/inference/v2/kernels/core_ops/test_post_ln.py

@@ -0,0 +1,47 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.core_ops import CUDAFPPostLN
+from ....v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_implementation(residual: torch.Tensor, hidden_states: torch.Tensor, gamma: torch.Tensor,
+                             beta: torch.Tensor, epsilon: float) -> torch.Tensor:
+    residual_f = residual.to(torch.float32)
+    hidden_states_f = hidden_states.to(torch.float32)
+    gamma_f = gamma.to(torch.float32)
+    beta_f = beta.to(torch.float32)
+    return torch.nn.functional.layer_norm(residual_f + hidden_states_f, (hidden_states_f.size(-1), ),
+                                          weight=gamma_f,
+                                          bias=beta_f,
+                                          eps=epsilon).to(hidden_states.dtype)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, channels", [(1, 4096), (37, 2048), (112, 14432), (1024, 6144)])
+@pytest.mark.parametrize("dtype", get_dtypes())
+def test_cuda_post_ln(tokens: int, channels: int, dtype: torch.dtype) -> None:
+
+    # Input vals
+    hidden_states = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    residual = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    gamma = torch.randn((channels), dtype=dtype, device=get_accelerator().current_device_name())
+    beta = torch.rand((channels), dtype=dtype, device=get_accelerator().current_device_name())
+    epsilon = 1e-5
+
+    # Reference output
+    ref_output = reference_implementation(residual, hidden_states, gamma, beta, epsilon)
+
+    # New output
+    post_ln_kernel = CUDAFPPostLN(hidden_states.size(-1), residual.dtype)
+    ds_output = torch.empty_like(residual)
+    post_ln_kernel(ds_output, residual, hidden_states, gamma, beta)
+
+    # Check
+    assert allclose(ds_output, ref_output)

tests/unit/inference/v2/kernels/core_ops/test_pre_ln.py

@@ -0,0 +1,51 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.core_ops import CUDAFPPreLN
+from ....v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_implementation(residual: torch.Tensor, hidden_states: torch.Tensor, gamma: torch.Tensor,
+                             beta: torch.Tensor, epsilon: float) -> torch.Tensor:
+    residual_f = residual.to(torch.float32)
+    hidden_states_f = hidden_states.to(torch.float32)
+    gamma_f = gamma.to(torch.float32)
+    beta_f = beta.to(torch.float32)
+    residual_out = residual_f + hidden_states_f
+    hidden_out = torch.nn.functional.layer_norm(residual_out, (hidden_states_f.size(-1), ),
+                                                weight=gamma_f,
+                                                bias=beta_f,
+                                                eps=epsilon)
+    return residual_out.to(hidden_states.dtype), hidden_out.to(hidden_states.dtype)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, channels", [(1, 4096), (37, 2048), (112, 14432), (1024, 6144)])
+@pytest.mark.parametrize("dtype", get_dtypes())
+def test_cuda_pre_ln(tokens: int, channels: int, dtype: torch.dtype) -> None:
+
+    # Input vals
+    hidden_states = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    residual = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    gamma = torch.randn((channels), dtype=dtype, device=get_accelerator().current_device_name())
+    beta = torch.rand((channels), dtype=dtype, device=get_accelerator().current_device_name())
+    epsilon = 1e-5
+
+    # Reference output
+    ref_output_res, ref_output_hid = reference_implementation(residual, hidden_states, gamma, beta, epsilon)
+
+    # New output
+    pre_ln_kernel = CUDAFPPreLN(hidden_states.size(-1), residual.dtype)
+    ds_output_res = torch.empty_like(residual)
+    ds_output_hid = torch.empty_like(hidden_states)
+    pre_ln_kernel(ds_output_res, ds_output_hid, residual, hidden_states, gamma, beta)
+
+    # Check
+    assert allclose(ds_output_res, ref_output_res)
+    assert allclose(ds_output_hid, ref_output_hid)

tests/unit/inference/v2/kernels/core_ops/test_rms_norm.py

@@ -0,0 +1,77 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import DtypeEnum
+from deepspeed.inference.v2.kernels.core_ops import CUDARMSNorm, CUDARMSPreNorm
+from ....v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_rms_norm(vals: torch.Tensor, gamma: torch.Tensor, epsilon: float = 1e-5) -> torch.Tensor:
+    variance = vals.to(torch.float32).pow(2).mean(-1, keepdim=True)
+    vals = vals * torch.rsqrt(variance + epsilon)
+
+    if gamma.dtype in [torch.float16, torch.bfloat16]:
+        vals = vals.to(gamma.dtype)
+
+    return gamma * vals
+
+
+def reference_rms_pre_norm(vals: torch.Tensor,
+                           residual: torch.Tensor,
+                           gamma: torch.Tensor,
+                           epsilon: float = 1e-5) -> torch.Tensor:
+    residual = residual + vals
+    return residual, reference_rms_norm(residual, gamma, epsilon)
+
+
+def _rms_norm_testing_helper(rows: int, channels: int, do_residual: bool, dtype: DtypeEnum) -> None:
+    device = get_accelerator().current_device_name()
+    t_dtype = dtype.value
+
+    vals = torch.randn((rows, channels), dtype=t_dtype, device=device)
+    gamma = torch.randn((channels), dtype=t_dtype, device=device)
+    epsilon = 1e-5
+
+    if do_residual:
+        residual_in = torch.randn((rows, channels), dtype=t_dtype, device=device)
+        ds_residual = residual_in.clone()
+
+        ref_residual, ref_output = reference_rms_pre_norm(vals, residual_in, gamma, epsilon)
+
+        kernel = CUDARMSPreNorm(channels, t_dtype, epsilon=epsilon)
+        ds_out = torch.empty_like(ds_residual)
+
+        kernel(ds_residual, ds_out, residual_in, vals, gamma)
+
+        assert allclose(ds_out, ref_output)
+        assert allclose(ds_residual, ref_residual)
+    else:
+
+        ref_output = reference_rms_norm(vals, gamma, epsilon)
+
+        kernel = CUDARMSNorm(channels, t_dtype, epsilon=epsilon)
+        ds_out = torch.empty_like(vals)
+
+        kernel(ds_out, vals, gamma)
+
+        assert allclose(ds_out, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("dtype", get_dtypes())
+@pytest.mark.parametrize("do_residual", [True, False])
+def test_rms_dtypes(dtype: DtypeEnum, do_residual: bool) -> None:
+    _rms_norm_testing_helper(883, 1024, do_residual, DtypeEnum(dtype))
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("rows, cols", [(1, 4096), (37, 2048), (112, 14432), (1024, 6144)])
+@pytest.mark.parametrize("do_residual", [True, False])
+def test_rms_shapes(rows: int, cols: int, do_residual: bool) -> None:
+    _rms_norm_testing_helper(rows, cols, do_residual, DtypeEnum.fp16)

tests/unit/inference/v2/kernels/cutlass_ops/__init__.py

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

tests/unit/inference/v2/kernels/cutlass_ops/test_moe_gemm.py

@@ -0,0 +1,113 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum
+from deepspeed.inference.v2.kernels.cutlass_ops import MoEGEMM
+from ....v2.inference_test_utils import allclose
+
+SINGLE_EXPERT_CASES = [(13, 2048, 2048), (256, 1024, 4096), (278, 5120, 2048), (893, 5120, 2560)]
+
+PYTORCH_ACT_FN_MAP = {
+    ActivationType.GELU: torch.nn.functional.gelu,
+    ActivationType.SILU: torch.nn.functional.silu,
+    ActivationType.RELU: torch.nn.functional.relu
+}
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, in_neurons, out_neurons", SINGLE_EXPERT_CASES)
+def test_single_expert(n_tokens: int, in_neurons: int, out_neurons: int) -> None:
+    """
+    Validate that the GEMM kernel produces identical results for a single GEMM instance.
+    """
+    device = get_accelerator().current_device()
+
+    activations = torch.rand((n_tokens, in_neurons), device=device, dtype=torch.float16) - 0.5
+    weights = torch.rand((1, in_neurons, out_neurons), device=device, dtype=torch.float16) - 0.5
+    biases = torch.randn((1, out_neurons), device=device, dtype=torch.float16)
+
+    weights_ref = weights.reshape(in_neurons, out_neurons)
+    biases_ref = biases.reshape(out_neurons)
+    ref_output = torch.matmul(activations, weights_ref) + biases_ref
+
+    moe_gemm = MoEGEMM(DtypeEnum.fp16, ActivationType.IDENTITY)
+    output = torch.empty((n_tokens, out_neurons), device=device, dtype=torch.float16)
+    cumsum_rows = torch.tensor([n_tokens], dtype=torch.int64, device=device)
+
+    moe_gemm(output, activations, weights, cumsum_rows, biases)
+    assert allclose(output, ref_output, tolerances=(1e-2, 1e-2))
+    get_accelerator().synchronize()
+
+
+def moe_test_helper(in_neurons: int, out_neurons: int, n_experts: int, max_tokens_per_expert: int,
+                    act_fn: ActivationType, dtype: DtypeEnum) -> None:
+    """
+    Helper function for validating the GEMM kernel for a single expert.
+    """
+    device = get_accelerator().current_device()
+
+    expert_allocations = torch.randint(0, max_tokens_per_expert, (n_experts, ), device=device, dtype=torch.int32)
+    cumsum_rows = expert_allocations.cumsum(dim=0)
+    print(cumsum_rows.dtype)
+
+    activations = torch.rand((cumsum_rows[-1], in_neurons), device=device, dtype=dtype.value) - 0.5
+    weights = torch.rand((n_experts, in_neurons, out_neurons), device=device, dtype=dtype.value) - 0.5
+    biases = torch.randn((n_experts, out_neurons), device=device, dtype=dtype.value)
+
+    out_ref = torch.empty((cumsum_rows[-1], out_neurons), device=device, dtype=dtype.value)
+
+    for expert_idx in range(n_experts):
+        start = cumsum_rows[expert_idx - 1] if expert_idx > 0 else 0
+        end = cumsum_rows[expert_idx]
+        activations_slice = activations[start:end]
+        weights_slice = weights[expert_idx]
+        biases_slice = biases[expert_idx]
+        out_ref[start:end] = torch.matmul(activations_slice, weights_slice) + biases_slice
+
+    if act_fn != ActivationType.IDENTITY:
+        act_fn_fn = PYTORCH_ACT_FN_MAP[act_fn]
+        out_ref = act_fn_fn(out_ref)
+
+    moe_gemm = MoEGEMM(DtypeEnum.fp16, act_fn)
+    output = torch.empty((cumsum_rows[-1], out_neurons), device=device, dtype=dtype.value)
+
+    moe_gemm(output, activations, weights, cumsum_rows, biases)
+
+    if dtype == DtypeEnum.bf16:
+        assert allclose(output, out_ref, tolerances=(1e-1, 1e-1))
+    else:
+        assert allclose(output, out_ref, tolerances=(1e-2, 1e-2))
+    get_accelerator().synchronize()
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("max_tokens_per_expert", [1, 4, 16, 64, 128])
+def test_multi_expert(max_tokens_per_expert: int) -> None:
+    """
+    Validate for multi-expert GEMM instances that the output is identical to the reference.
+    """
+    moe_test_helper(5120, 2048, 64, max_tokens_per_expert, ActivationType.IDENTITY, DtypeEnum.fp16)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("act_fn", [ActivationType.GELU, ActivationType.SILU, ActivationType.RELU])
+def test_act_fns(act_fn: ActivationType) -> None:
+    """
+    Validate activation function behavior.
+    """
+    moe_test_helper(5120, 2048, 64, 32, act_fn, DtypeEnum.fp16)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("dtype", [DtypeEnum.fp16, DtypeEnum.bf16])
+def test_dtypes(dtype: DtypeEnum) -> None:
+    """
+    Validate data type behavior.
+    """
+    moe_test_helper(5120, 2048, 64, 32, ActivationType.IDENTITY, dtype)

tests/unit/inference/v2/kernels/ragged_ops/__init__.py

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

tests/unit/inference/v2/kernels/ragged_ops/ragged_testing_utils.py

@@ -0,0 +1,300 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import random
+from typing import List, Optional, Tuple
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.ragged import (
+    AllocationMode,
+    DSSequenceDescriptor,
+    DSStateManager,
+    DSStateManagerConfig,
+    KVCacheConfig,
+    MemoryConfig,
+    PlaceholderSequenceDescriptor,
+    RaggedBatchWrapper,
+)
+from ....v2.inference_test_utils import allclose
+
+
+def build_simple_batch(seq_lens: List[int],
+                       vocab_range: Optional[int] = 100,
+                       padding: Optional[bool] = False) -> RaggedBatchWrapper:
+    """
+    Construct a simple batch with the given sequence lengths. This method should not
+    be used for for testing scenarios that require information about KV or sequence
+    history.
+    """
+    total_tokens = max(sum(seq_lens), 1024)
+    n_seqs = max(len(seq_lens), 128)
+
+    config = DSStateManagerConfig(max_tracked_sequences=n_seqs,
+                                  max_ragged_sequence_count=n_seqs,
+                                  max_ragged_batch_size=total_tokens)
+    batch = RaggedBatchWrapper(config)
+
+    batch.clear()
+
+    for seq_len in seq_lens:
+        seq_desc = PlaceholderSequenceDescriptor()
+        tokens = torch.randint(0, vocab_range, (seq_len, ))
+        batch.insert_sequence(seq_desc, tokens)
+
+    batch.finalize(padding=padding)
+
+    return batch
+
+
+def build_complex_batch(seq_params: List[Tuple[int, int, int]],
+                        kv_block_size: int,
+                        vocab_range: Optional[int] = 100,
+                        padding: Optional[bool] = False) -> Tuple[RaggedBatchWrapper, int]:
+    """
+    Construct a fully paramtrized batch with the given sequence lengths. This method
+    can be used to construct more realistic inputs for testing scenarios that will interact
+    with all the members of the RaggedBatchWrapper.
+    """
+    seq_lens = [seq_param[0] for seq_param in seq_params]
+    total_tokens = max(sum(seq_lens), 1024)
+    n_seqs = max(len(seq_lens), 128)
+
+    config = DSStateManagerConfig(max_tracked_sequences=n_seqs,
+                                  max_ragged_sequence_count=n_seqs,
+                                  max_ragged_batch_size=total_tokens)
+    batch = RaggedBatchWrapper(config)
+
+    batch.clear()
+
+    total_kv_blocks = 0
+
+    for seq_len, n_seen_tokens, kv_ptr in seq_params:
+        n_kv_blocks = (seq_len + n_seen_tokens + kv_block_size - 1) // kv_block_size
+        seq_desc = PlaceholderSequenceDescriptor(seen_tokens=n_seen_tokens,
+                                                 cur_allocated_blocks=n_kv_blocks,
+                                                 kv_blocks_ptr=kv_ptr)
+        tokens = torch.randint(0, vocab_range, (seq_len, ))
+        batch.insert_sequence(seq_desc, tokens)
+        total_kv_blocks += n_kv_blocks
+
+    batch.finalize(padding=padding)
+
+    return batch, total_kv_blocks
+
+
+def build_batch_and_manager(
+    seq_params: List[Tuple[int, int]],
+    head_size: int,
+    n_heads_kv: int,
+    kv_block_size: int,
+    vocab_range: Optional[int] = 100,
+    padding: Optional[bool] = False,
+    kv_fill: Optional[List[torch.Tensor]] = None
+) -> Tuple[RaggedBatchWrapper, DSStateManager, List[DSSequenceDescriptor]]:
+    """
+    Will construct and populate a batch and KVCache with the given sequence parameters.
+
+    Arguments:
+        seq_params (List[Tuple[int, int]]): A list of tuples containing the sequence length and
+            the number of tokens that have already been seen for that sequence.
+        head_size (int): The size of each attention head.
+        n_heads_kv (int): The number of attention heads for the KV-cache.
+        kv_block_size (int): The size of each block in the KV-cache.
+        vocab_range (Optional[int]): The range of the vocabulary. Defaults to 100.
+        padding (Optional[bool]): Whether to pad the batch. Defaults to False.
+        kv_fill (Optional[List[torch.Tensor]]): A list of tensors to use to populate the KV-cache.
+            If this is not provided, the KV-cache will be treated as empty and the contents should
+            not be relied upon. NOTE(cmikeh2): This functionality relies on the functionality
+            of LinearBlockedKVCopy. If tests relying on this feature are failing, make sure that
+            LinearBlockedKVCopy is working correctly.
+    """
+    seq_lens = [seq_param[0] for seq_param in seq_params]
+    fill_lens = [seq_param[1] for seq_param in seq_params]
+    max_created_batch_len = max(sum(seq_lens), sum(fill_lens))
+    total_tokens = max(max_created_batch_len, 1024)
+    n_seqs = max(len(seq_lens), 128)
+
+    req_kv_blocks = [None] * n_seqs
+    total_kv_blocks = 0
+    for i, (seq_len, n_seen_tokens) in enumerate(seq_params):
+        req_kv_blocks[i] = (seq_len + n_seen_tokens + kv_block_size - 1) // kv_block_size
+        total_kv_blocks += req_kv_blocks[i]
+
+    kv_config = KVCacheConfig(block_size=kv_block_size,
+                              num_allocation_groups=1,
+                              cache_shape=(1, n_heads_kv, head_size))
+    memory_config = MemoryConfig(mode=AllocationMode.ALLOCATE, size=total_kv_blocks)
+
+    config = DSStateManagerConfig(max_tracked_sequences=n_seqs,
+                                  max_ragged_sequence_count=n_seqs,
+                                  max_ragged_batch_size=total_tokens,
+                                  memory_config=memory_config)
+
+    batch = RaggedBatchWrapper(config)
+    state_manager = DSStateManager(config, (kv_config, ))
+
+    # At the beginning of operation, the design of the allocator is such that it will return
+    # linear blocks of memory. The following will "warm up" the allocator so that we can be
+    # more certain that code is not dependent on this behavior.
+    all_allocs = []
+    for _ in range(20):
+        decision = random.randint(0, 1)
+
+        if decision == 0:
+            blocks_to_allocate = random.randint(0, total_kv_blocks)
+            if blocks_to_allocate <= state_manager.free_blocks[0] and blocks_to_allocate > 0:
+                all_allocs.append(state_manager.allocate_blocks(blocks_to_allocate))
+        else:
+            if len(all_allocs) > 0:
+                idx = random.randint(0, len(all_allocs) - 1)
+                state_manager._kv_cache.free(all_allocs[idx])
+
+                del all_allocs[idx]
+
+    for alloc in all_allocs:
+        state_manager._kv_cache.free(alloc)
+
+    assert state_manager.free_blocks[0] == total_kv_blocks
+
+    batch.clear()
+    seq_descs = []
+
+    if kv_fill is None or sum(fill_lens) == 0:
+        for i, (seq_len, n_seen_tokens) in enumerate(seq_params):
+            # Create empty descriptor
+            seq_desc = state_manager.get_or_create_sequence(i)
+
+            # Update `seen_tokens` in the descriptor
+            seq_desc.pre_forward(n_seen_tokens)
+            seq_desc.post_forward()
+
+            # Ensure there's enough KV-cache for the sequence
+            kv_block_ids = state_manager.allocate_blocks(req_kv_blocks[i])
+            print(f"Allocated {req_kv_blocks[i]} blocks for sequence {i}: {kv_block_ids}")
+            seq_desc.extend_kv_cache(kv_block_ids)
+
+            # Insert sequence into batch
+            tokens = torch.randint(0, vocab_range, (seq_len, ))
+            batch.insert_sequence(seq_desc, tokens)
+            seq_desc.pre_forward(seq_len)
+            seq_descs.append(seq_desc)
+    else:
+        qkv = torch.empty((total_tokens, (n_heads_kv * 3) * head_size),
+                          dtype=torch.float16,
+                          device=get_accelerator().current_device())
+        fills_as_tensor = torch.tensor(fill_lens, dtype=torch.int32)
+        fill_cumsum = torch.cat((torch.tensor([0], dtype=torch.int32), torch.cumsum(fills_as_tensor, dim=0)))
+
+        for i, (_, n_seen_tokens) in enumerate(seq_params):
+            # Create empty descriptor
+            seq_desc = state_manager.get_or_create_sequence(i)
+
+            # Update `seen_tokens` in the descriptor
+            if n_seen_tokens > 0:
+                dummy_fill_toks = torch.randint(0, vocab_range, (n_seen_tokens, ))
+                batch.insert_sequence(seq_desc, dummy_fill_toks)
+                seq_desc.pre_forward(n_seen_tokens)
+
+            # Ensure there's enough KV-cache for the sequence
+            kv_block_ids = state_manager.allocate_blocks(req_kv_blocks[i])
+            print(f"Allocated {req_kv_blocks[i]} blocks for sequence {i}: {kv_block_ids}")
+            seq_desc.extend_kv_cache(kv_block_ids)
+            seq_descs.append(seq_desc)
+
+            if n_seen_tokens == 0:
+                continue
+
+            assert kv_fill[i].shape[0] == n_seen_tokens
+            assert kv_fill[i].shape[1] == n_heads_kv * head_size * 2
+
+            local_q = torch.randn((n_seen_tokens, n_heads_kv * head_size), dtype=torch.float16, device=qkv.device)
+            local_qkv = torch.cat((local_q, kv_fill[i]), dim=1)
+            qkv[fill_cumsum[i]:fill_cumsum[i + 1]] = local_qkv
+
+        batch.finalize(padding=padding)
+
+        from deepspeed.inference.v2.kernels.ragged_ops import LinearBlockedKVCopy
+        kv_copy = LinearBlockedKVCopy(head_size, n_heads_kv, n_heads_kv, torch.float16)
+        kv_cache = state_manager.get_cache(0)
+        kv_copy(kv_cache, qkv, batch)
+
+        for seq_desc in seq_descs:
+            if seq_desc.in_flight_tokens > 0:
+                seq_desc.post_forward()
+
+        batch.clear()
+
+        for i, (seq_len, _) in enumerate(seq_params):
+            seq_desc = state_manager.get_or_create_sequence(i)
+            tokens = torch.randint(0, vocab_range, (seq_len, ))
+            batch.insert_sequence(seq_desc, tokens)
+            seq_desc.pre_forward(seq_len)
+
+            # We will skip KV cache allocation here because we did a lump allocation above
+            # for both the fill and the sequence itself.
+
+    batch.finalize(padding=padding)
+
+    return batch, state_manager, seq_descs
+
+
+def validate_kv_cache(kv_cache: torch.Tensor,
+                      k: torch.Tensor,
+                      v: torch.Tensor,
+                      seq_descs: List[DSSequenceDescriptor],
+                      batch: RaggedBatchWrapper,
+                      exact: bool = True) -> None:
+    """
+    Given a QKV tensor and a KV cache, validate that the cache contains the correct values.
+    """
+    block_size = kv_cache.shape[1]
+    n_kv_heads = kv_cache.shape[3]
+    head_size = kv_cache.shape[4]
+
+    inflight_descs = batch.inflight_seq_descriptors(on_device=False)[:batch.current_sequences]
+
+    if inflight_descs.shape[0] != len(seq_descs):
+        raise ValueError("The number of sequence descriptors does not match the number of sequences in the batch.")
+
+    for seq_desc, inflight_seq in zip(seq_descs, inflight_descs):
+        start_idx = inflight_seq[0]
+        assigned_kv_blocks = seq_desc.kv_cache_ids(on_device=False)
+
+        real_k_values = k[start_idx:start_idx + seq_desc.in_flight_tokens]
+        real_v_values = v[start_idx:start_idx + seq_desc.in_flight_tokens]
+
+        start_block_idx = seq_desc.seen_tokens // block_size
+        local_start_idx = 0
+        cur_start_idx = seq_desc.seen_tokens
+
+        for block_idx in range(start_block_idx, seq_desc.cur_allocated_blocks):
+            block = kv_cache[assigned_kv_blocks[0, block_idx].item()]
+            block_start_idx = cur_start_idx % block_size
+            n_tokens_to_check = min(block_size - block_start_idx, seq_desc.in_flight_tokens - local_start_idx)
+            block_end_idx = block_start_idx + n_tokens_to_check
+
+            if exact:
+                assert torch.equal(
+                    block[block_start_idx:block_end_idx, 0, :, :],
+                    real_k_values[local_start_idx:local_start_idx + n_tokens_to_check].reshape(
+                        n_tokens_to_check, n_kv_heads, head_size))
+                assert torch.equal(
+                    block[block_start_idx:block_end_idx, 1, :, :],
+                    real_v_values[local_start_idx:local_start_idx + n_tokens_to_check].reshape(
+                        n_tokens_to_check, n_kv_heads, head_size))
+            else:
+                assert allclose(
+                    block[block_start_idx:block_end_idx, 0, :, :],
+                    real_k_values[local_start_idx:local_start_idx + n_tokens_to_check].reshape(
+                        n_tokens_to_check, n_kv_heads, head_size))
+                assert allclose(
+                    block[block_start_idx:block_end_idx, 1, :, :],
+                    real_v_values[local_start_idx:local_start_idx + n_tokens_to_check].reshape(
+                        n_tokens_to_check, n_kv_heads, head_size))
+
+            local_start_idx += n_tokens_to_check
+            cur_start_idx += n_tokens_to_check

tests/unit/inference/v2/kernels/ragged_ops/test_atom_builder.py

@@ -0,0 +1,45 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.inference.v2.kernels.ragged_ops import AtomBuilder
+from .ragged_testing_utils import build_complex_batch
+
+Q_BLOCK_SIZE = 128
+KV_BLOCK_SIZE = 128
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('seq_params', [(1, 0, 0), (1, 228, 0), (383, 0, 0), (1, 494, 0)])
+def test_single_sequence(seq_params) -> None:
+    seq_len, n_seen_tokens, _ = seq_params
+
+    batch, _ = build_complex_batch([seq_params], kv_block_size=KV_BLOCK_SIZE, padding=False)
+    atom_builder = AtomBuilder()
+
+    atoms = torch.empty((8, 8), dtype=torch.int32, device=torch.device("cpu"))
+    atoms, n_atoms = atom_builder(atoms, batch, Q_BLOCK_SIZE, KV_BLOCK_SIZE)
+
+    calc_n_atoms = (seq_len + 127) // 128
+
+    assert n_atoms == calc_n_atoms
+
+    for i, atom in enumerate(atoms[:n_atoms]):
+        # Since the ptr was 0, first 2 elements should be 0
+        assert atom[0] == 0
+        assert atom[1] == 0
+
+        # Since we have a single sequence, the q_start_idx should always be
+        # whichever atom we're on multiplied by the block size
+        assert atom[2] == i * Q_BLOCK_SIZE
+        assert atom[3] == min(Q_BLOCK_SIZE, seq_len - i * Q_BLOCK_SIZE)
+        total_toks = i * Q_BLOCK_SIZE + min(Q_BLOCK_SIZE, seq_len - i * Q_BLOCK_SIZE)
+
+        assert atom[4] == (total_toks + n_seen_tokens + KV_BLOCK_SIZE - 1) // KV_BLOCK_SIZE
+        assert atom[5] == (total_toks + n_seen_tokens)
+
+        assert atom[6] == n_seen_tokens + i * Q_BLOCK_SIZE

tests/unit/inference/v2/kernels/ragged_ops/test_blocked_flash.py

@@ -0,0 +1,197 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import itertools
+
+from typing import List, Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import DtypeEnum
+from deepspeed.inference.v2.kernels.ragged_ops import (
+    AtomBuilder,
+    BlockedFlashAttn,
+    get_q_block_size,
+    get_kv_block_size,
+    LinearBlockedKVCopy,
+)
+from deepspeed.inference.v2.ragged import split_kv
+from deepspeed.ops.op_builder import RaggedUtilsBuilder
+
+from .ragged_testing_utils import build_batch_and_manager
+from ....v2.inference_test_utils import allclose
+
+try:
+    from flash_attn.flash_attn_interface import flash_attn_varlen_func
+    validate_accuracy = True
+except ImportError:
+    validate_accuracy = False
+"""
+NOTE(cmikeh2): These tests depend on atom construction and KV-cache copying to behave correctly.
+If one or the other of those is not working, then these tests will fail. Before debugging here,
+make sure that the atom construction and KV-cache copying tests are passing.
+"""
+
+
+def _blocked_flash_testing_helper(head_size: int, n_heads_q: int, n_heads_kv: int,
+                                  seq_params: List[Tuple[int, int]]) -> None:
+    """
+    Helper function for testing blocked flash attention. Used to enable parametrize to only set up
+    a subset of parameters before being passed to the unified test function.
+    """
+    q_block_size = get_q_block_size(head_size)
+    kv_block_size = get_kv_block_size(head_size)
+
+    kvs = []
+    for _, history_len in seq_params:
+        if history_len > 0:
+            kvs.append(
+                torch.randn((history_len, 2 * n_heads_kv * head_size),
+                            device=get_accelerator().current_device(),
+                            dtype=torch.float16))
+        else:
+            kvs.append(None)
+
+    batch, state_manager, _ = build_batch_and_manager(seq_params, head_size, n_heads_kv, kv_block_size, kv_fill=kvs)
+
+    atom_builder = AtomBuilder()
+    kv_copy = LinearBlockedKVCopy(head_size, n_heads_q, n_heads_kv, DtypeEnum.fp16)
+    atom_flash = BlockedFlashAttn(head_size, DtypeEnum.fp16)
+
+    total_atoms = sum((seq[0] + q_block_size - 1) // q_block_size for seq in seq_params)
+    atoms = torch.empty((total_atoms, 8), dtype=torch.int32, device=get_accelerator().current_device())
+    alloc_func = RaggedUtilsBuilder().load().allocate_fast_host_buffer
+    atoms_host = alloc_func(atoms)
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=get_accelerator().current_device(),
+                      dtype=torch.float16)
+
+    atoms_host, n_atoms = atom_builder(atoms_host, batch, q_block_size, kv_block_size)
+    atoms.copy_(atoms_host[:n_atoms])
+
+    kv_cache = state_manager.get_cache(0)
+    kv_copy(kv_cache, qkv, batch)
+
+    out = torch.empty((batch.current_tokens, head_size * n_heads_q),
+                      device=get_accelerator().current_device(),
+                      dtype=torch.float16)
+    k_cache, v_cache = split_kv(kv_cache)
+    q = qkv[:, :head_size * n_heads_q]
+
+    atom_flash(out, q, k_cache, v_cache, atoms, 1.0)
+
+    if validate_accuracy:
+        cu_seqlens_q = torch.tensor([0] + list(itertools.accumulate([seq[0] for seq in seq_params])),
+                                    dtype=torch.int32,
+                                    device=get_accelerator().current_device())
+        cu_seqlens_kv = torch.tensor([0] + list(itertools.accumulate([seq[1] + seq[0] for seq in seq_params])),
+                                     dtype=torch.int32,
+                                     device=get_accelerator().current_device())
+
+        inflight_kv = qkv[:, head_size * n_heads_q:]
+        full_kvs = []
+        for i, kv in enumerate(kvs):
+            if kv is not None:
+                full_kvs.append(torch.cat([kv, inflight_kv[cu_seqlens_q[i]:cu_seqlens_q[i + 1]]], dim=0))
+            else:
+                full_kvs.append(inflight_kv[cu_seqlens_q[i]:cu_seqlens_q[i + 1]])
+        run_kvs = torch.cat(full_kvs, dim=0)
+        k = run_kvs[:, :head_size * n_heads_kv]
+        v = run_kvs[:, head_size * n_heads_kv:]
+
+        q_ref = q.reshape((batch.current_tokens, n_heads_q, head_size))
+        k_ref = k.reshape((k.shape[0], n_heads_kv, head_size))
+        v_ref = v.reshape((v.shape[0], n_heads_kv, head_size))
+
+        max_seqlen_q = max([seq[0] for seq in seq_params])
+        max_seqlen_kv = max([seq[1] + seq[0] for seq in seq_params])
+
+        ref_o = flash_attn_varlen_func(q_ref,
+                                       k_ref,
+                                       v_ref,
+                                       cu_seqlens_q,
+                                       cu_seqlens_kv,
+                                       max_seqlen_q,
+                                       max_seqlen_kv,
+                                       softmax_scale=1.0,
+                                       causal=True)
+
+        ref_o = ref_o.reshape(batch.current_tokens, head_size * n_heads_q)
+
+        assert allclose(out, ref_o)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens", [2, 33, 65, 128, 256, 2037])
+def test_single_prompt(n_tokens: int) -> None:
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(n_tokens, 0)]
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("prompt_lengths", [(128, 128), (192, 38), (514, 713), (83, 312, 610)])
+def test_multiple_prompts(prompt_lengths: Tuple[int, int]) -> None:
+    """
+    Test multiple prompts in a single batch.
+    """
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(prompt_lengths[i], 0) for i in range(len(prompt_lengths))]
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("seq_params", [(1, 34), (43, 40), (1, 144), (64, 128), (332, 628)])
+def test_continuation(seq_params: Tuple[int, int]) -> None:
+    """
+    Test continued generation/prompt processing.
+    """
+    head_size = 64
+    n_heads_q = 32
+    n_heads_kv = 32
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, [seq_params])
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("head_size", [64, 128])
+def test_head_size(head_size: int) -> None:
+    n_heads_q = 16
+    n_heads_kv = 16
+    seq_params = [(128, 128), (192, 38), (1, 814)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("head_config", [(32, 8), (64, 16), (40, 8)])
+def test_gqa(head_config: Tuple[int, int]) -> None:
+    head_size = 128
+    n_heads_q = head_config[0]
+    n_heads_kv = head_config[1]
+
+    seq_params = [(128, 128), (192, 38), (1, 814)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+def test_fully_composed() -> None:
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(332, 628), (1, 718), (1, 323), (180, 5), (224, 0)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)

tests/unit/inference/v2/kernels/ragged_ops/test_blocked_kv_copy.py

@@ -0,0 +1,112 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.ragged_ops import LinearBlockedKVCopy
+from .ragged_testing_utils import build_batch_and_manager, validate_kv_cache
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, history_size", [(1, 0), (17, 0), (33, 8), (63, 1)])
+@pytest.mark.parametrize("head_size", [64, 80, 96, 128])
+def test_single_sequence_single_block(n_tokens: int, history_size: int, head_size: int):
+    """
+    Validate that the copy works correctly
+    """
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch, state_manager, seq_descs = build_batch_and_manager([(n_tokens, history_size)], head_size, n_heads_kv,
+                                                              kv_block_size)
+
+    assert batch.current_sequences == 1
+    assert batch.current_tokens == n_tokens
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+
+    kv_cache = state_manager.get_cache(0)
+
+    copy_impl = LinearBlockedKVCopy(head_size, n_heads_q, n_heads_kv, torch.float16)
+    copy_impl(kv_cache, qkv, batch)
+
+    k = qkv[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv[:, head_size * (n_heads_q + n_heads_kv):]
+
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, history_size", [(128, 0), (177, 0), (169, 8), (117, 88)])
+@pytest.mark.parametrize("head_size", [64, 80, 96, 128])
+def test_single_sequence_multiple_blocks(n_tokens: int, history_size: int, head_size: int):
+    """
+    Validate that the copy works correctly
+    """
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch, state_manager, seq_descs = build_batch_and_manager([(n_tokens, history_size)], head_size, n_heads_kv,
+                                                              kv_block_size)
+
+    assert batch.current_sequences == 1
+    assert batch.current_tokens == n_tokens
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+
+    kv_cache = state_manager.get_cache(0)
+
+    copy_impl = LinearBlockedKVCopy(head_size, n_heads_q, n_heads_kv, torch.float16)
+    copy_impl(kv_cache, qkv, batch)
+
+    k = qkv[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv[:, head_size * (n_heads_q + n_heads_kv):]
+
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("head_size", [64, 80, 96, 128])
+def test_multi_sequence(head_size: int) -> None:
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch_config = [
+        (128, 0),
+        (177, 0),
+        (169, 8),
+        (117, 88),
+        (1, 293),
+        (1, 733),
+        (1, 33),
+    ]
+
+    batch, state_manager, seq_descs = build_batch_and_manager(batch_config, head_size, n_heads_kv, kv_block_size)
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+
+    kv_cache = state_manager.get_cache(0)
+
+    copy_impl = LinearBlockedKVCopy(head_size, n_heads_q, n_heads_kv, torch.float16)
+    copy_impl(kv_cache, qkv, batch)
+
+    k = qkv[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv[:, head_size * (n_heads_q + n_heads_kv):]
+
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch)

tests/unit/inference/v2/kernels/ragged_ops/test_blocked_rotary_emb.py

@@ -0,0 +1,258 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import List, Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.ragged_ops import BlockedRotaryEmbeddings, BlockedTrainedRotaryEmbeddings
+from deepspeed.inference.v2.ragged import RaggedBatchWrapper, DSSequenceDescriptor
+from .ragged_testing_utils import build_batch_and_manager, validate_kv_cache
+from ....v2.inference_test_utils import allclose
+"""
+NOTE(cmikeh2): It is very possible to see unit test failures (even on FP16) depending on when
+certain values are casted up to or down from float32. If we are seeing accuracy issues, we should
+make sure we are aligning on the training implementation's cast pattern here, given these tolerances
+tend to be sufficient elsewhere.
+"""
+
+
+def rotary_pos_embs(q: torch.Tensor,
+                    k: torch.Tensor,
+                    seq_descs: List[DSSequenceDescriptor],
+                    batch: RaggedBatchWrapper,
+                    head_size: int,
+                    rotary_dim: int = -1) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    rotary_dim = rotary_dim if rotary_dim >= 0 else head_size
+
+    def make_cos_sin_emb(seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        t = torch.arange(seq_len, dtype=torch.float32, device=get_accelerator().current_device())
+        inv_freq = (1.0 / (10000.0**(torch.arange(
+            0, rotary_dim, 2, dtype=torch.float32, device=get_accelerator().current_device()) / rotary_dim))).half()
+
+        freqs = torch.einsum("i,j->ij", t, inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+
+        return torch.cos(emb)[:, None, :], torch.sin(emb)[:, None, :], inv_freq
+
+    def rotate_half(x: torch.Tensor) -> torch.Tensor:
+        return torch.cat((-x[..., x.shape[-1] // 2:], x[..., :x.shape[-1] // 2]), dim=-1)
+
+    cos, sin, freqs = make_cos_sin_emb(1024)
+
+    q_out = torch.empty_like(q)
+    k_out = torch.empty_like(k)
+    n_heads_q = q.shape[1] // head_size
+    n_heads_kv = k.shape[1] // head_size
+
+    inflight_descs = batch.inflight_seq_descriptors(on_device=False)[:batch.current_sequences]
+
+    if inflight_descs.shape[0] != len(seq_descs):
+        raise ValueError("The number of sequence descriptors does not match the number of sequences in the batch.")
+
+    for seq_desc, inflight_seq in zip(seq_descs, inflight_descs):
+        start_idx = inflight_seq[0]
+        n_tokens = seq_desc.in_flight_tokens
+
+        q_src = q[start_idx:start_idx + n_tokens].reshape(n_tokens, n_heads_q, head_size).float()
+        k_src = k[start_idx:start_idx + n_tokens].reshape(n_tokens, n_heads_kv, head_size).float()
+        freq_start_offset = seq_desc.seen_tokens
+
+        q_src_rot = q_src[:, :, :rotary_dim]
+        k_src_rot = k_src[:, :, :rotary_dim]
+
+        cos_chunk = cos[range(freq_start_offset, freq_start_offset + n_tokens)]
+        sin_chunk = sin[range(freq_start_offset, freq_start_offset + n_tokens)]
+
+        q_rot = q_src_rot * cos_chunk + rotate_half(q_src_rot) * sin_chunk
+        k_rot = k_src_rot * cos_chunk + rotate_half(k_src_rot) * sin_chunk
+
+        q_emb = torch.cat((q_rot, q_src[:, :, rotary_dim:]), dim=-1)
+        k_emb = torch.cat((k_rot, k_src[:, :, rotary_dim:]), dim=-1)
+
+        q_out[start_idx:start_idx + n_tokens] = q_emb.reshape(n_tokens, n_heads_q * head_size).to(q_out.dtype)
+        k_out[start_idx:start_idx + n_tokens] = k_emb.reshape(n_tokens, n_heads_kv * head_size).to(k_out.dtype)
+
+    return q_out, k_out, freqs
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, history_size", [(1, 0), (17, 0), (33, 15), (1, 63)])
+@pytest.mark.parametrize("trained_emb", [False, True])
+@pytest.mark.parametrize("head_size", [64, 80, 96])
+def test_single_sequence_single_block(n_tokens: int, history_size: int, trained_emb: bool, head_size: int):
+    """
+    Validate that the copy works correctly
+    """
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch, state_manager, seq_descs = build_batch_and_manager([(n_tokens, history_size)], head_size, n_heads_kv,
+                                                              kv_block_size)
+
+    assert batch.current_sequences == 1
+    assert batch.current_tokens == n_tokens
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+    qkv_ref = qkv.clone()
+
+    q = qkv_ref[:, :head_size * n_heads_q]
+    k = qkv_ref[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv_ref[:, head_size * (n_heads_q + n_heads_kv):]
+
+    q_ref, k, freqs = rotary_pos_embs(q, k, seq_descs, batch, head_size)
+    freqs = freqs.half()
+
+    kv_cache = state_manager.get_cache(0)
+
+    if trained_emb:
+        copy_impl = BlockedTrainedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16)
+        copy_impl(kv_cache, qkv, batch, freqs)
+    else:
+        copy_impl = BlockedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16, head_size, 10000.0)
+        copy_impl(kv_cache, qkv, batch)
+
+    assert allclose(qkv[:, :head_size * n_heads_q], q_ref)
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch, exact=False)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, history_size", [(128, 0), (177, 0), (169, 8), (117, 88)])
+@pytest.mark.parametrize("trained_emb", [False, True])
+@pytest.mark.parametrize("head_size", [64, 80, 96])
+def test_single_sequence_multiple_blocks(n_tokens: int, history_size: int, trained_emb: bool, head_size: int):
+    """
+    Validate that the copy works correctly
+    """
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch, state_manager, seq_descs = build_batch_and_manager([(n_tokens, history_size)], head_size, n_heads_kv,
+                                                              kv_block_size)
+
+    assert batch.current_sequences == 1
+    assert batch.current_tokens == n_tokens
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+    qkv_ref = qkv.clone()
+
+    q = qkv_ref[:, :head_size * n_heads_q]
+    k = qkv_ref[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv_ref[:, head_size * (n_heads_q + n_heads_kv):]
+
+    q_ref, k, freqs = rotary_pos_embs(q, k, seq_descs, batch, head_size)
+    freqs = freqs.half()
+
+    kv_cache = state_manager.get_cache(0)
+
+    if trained_emb:
+        copy_impl = BlockedTrainedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16)
+        copy_impl(kv_cache, qkv, batch, freqs)
+    else:
+        copy_impl = BlockedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16, head_size, 10000.0)
+        copy_impl(kv_cache, qkv, batch)
+
+    assert allclose(qkv[:, :head_size * n_heads_q], q_ref)
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch, exact=False)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("trained_emb", [False, True])
+@pytest.mark.parametrize("head_size", [64, 80, 96])
+def test_multi_sequences(trained_emb: bool, head_size: int) -> None:
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch_config = [
+        (128, 0),
+        (177, 0),
+        (169, 8),
+        (117, 88),
+        (1, 293),
+        (1, 733),
+        (1, 33),
+    ]
+
+    batch, state_manager, seq_descs = build_batch_and_manager(batch_config, head_size, n_heads_kv, kv_block_size)
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+    qkv_ref = qkv.clone()
+
+    q = qkv_ref[:, :head_size * n_heads_q]
+    k = qkv_ref[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv_ref[:, head_size * (n_heads_q + n_heads_kv):]
+
+    q_ref, k, freqs = rotary_pos_embs(q, k, seq_descs, batch, head_size)
+    freqs = freqs.half()
+
+    kv_cache = state_manager.get_cache(0)
+
+    if trained_emb:
+        copy_impl = BlockedTrainedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16)
+        copy_impl(kv_cache, qkv, batch, freqs)
+    else:
+        copy_impl = BlockedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16, head_size, 10000.0)
+        copy_impl(kv_cache, qkv, batch)
+
+    assert allclose(qkv[:, :head_size * n_heads_q], q_ref)
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch, exact=False)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("head_size", [80, 96])
+def test_rotary_dim(head_size: int) -> None:
+    trained_emb = False
+    rotary_dim = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+    kv_block_size = 64
+    device = get_accelerator().current_device()
+
+    batch_config = [
+        (128, 0),
+        (177, 0),
+        (169, 8),
+        (117, 88),
+        (1, 293),
+        (1, 733),
+        (1, 33),
+    ]
+
+    batch, state_manager, seq_descs = build_batch_and_manager(batch_config, head_size, n_heads_kv, kv_block_size)
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=device,
+                      dtype=torch.float16)
+    qkv_ref = qkv.clone()
+
+    q = qkv_ref[:, :head_size * n_heads_q]
+    k = qkv_ref[:, head_size * n_heads_q:head_size * (n_heads_q + n_heads_kv)]
+    v = qkv_ref[:, head_size * (n_heads_q + n_heads_kv):]
+
+    q_ref, k, freqs = rotary_pos_embs(q, k, seq_descs, batch, head_size, rotary_dim=rotary_dim)
+    freqs = freqs.half()
+
+    kv_cache = state_manager.get_cache(0)
+
+    copy_impl = BlockedRotaryEmbeddings(head_size, n_heads_q, n_heads_kv, torch.float16, rotary_dim, 10000.0)
+    copy_impl(kv_cache, qkv, batch)
+
+    assert allclose(qkv[:, :head_size * n_heads_q], q_ref)
+    validate_kv_cache(kv_cache, k, v, seq_descs, batch, exact=False)

tests/unit/inference/v2/kernels/ragged_ops/test_logits_gather.py

@@ -0,0 +1,96 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import List
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.ragged_ops import RaggedLogitsGather
+from ....v2.inference_test_utils import allclose, get_dtypes
+from .ragged_testing_utils import build_simple_batch
+
+
+def baseline_implementation(hidden_states: torch.Tensor, seq_lens: List[int]) -> torch.Tensor:
+    output = torch.empty((len(seq_lens), hidden_states.shape[1]),
+                         dtype=hidden_states.dtype,
+                         device=hidden_states.device)
+
+    offset = 0
+    for i, seq_len in enumerate(seq_lens):
+        output[i] = hidden_states[offset + seq_len - 1]
+        offset += seq_len
+
+    return output
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('dtype', get_dtypes())
+def test_supported_dtypes(dtype: torch.dtype) -> None:
+    """
+    Validate support on nominally supported data types.
+    """
+    model_dim = 4096
+
+    batch = build_simple_batch([256], padding=False)
+    hidden_states = torch.randn((batch.current_tokens, model_dim),
+                                dtype=dtype,
+                                device=get_accelerator().current_device())
+
+    reference_result = baseline_implementation(hidden_states, [256])
+
+    kernel = RaggedLogitsGather(model_dim, dtype)
+    output = torch.empty_like(reference_result)
+    kernel(output, hidden_states, batch)
+
+    assert allclose(output, reference_result)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('seq_lens', [[128, 64, 192, 32], [57, 112, 63, 89, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1],
+                                      [63, 27, 74, 83, 32, 17, 1, 1, 1, 1, 1]])
+def test_multiple_sequences(seq_lens: List[int]) -> None:
+    """
+    Validate support on more multi-sequence inputs.
+    """
+    model_dim = 4096
+    dtype = torch.float16
+
+    batch = build_simple_batch(seq_lens, padding=False)
+    hidden_states = torch.randn((batch.current_tokens, model_dim),
+                                dtype=dtype,
+                                device=get_accelerator().current_device())
+
+    reference_result = baseline_implementation(hidden_states, seq_lens)
+
+    kernel = RaggedLogitsGather(model_dim, dtype)
+    output = torch.empty_like(reference_result)
+    kernel(output, hidden_states, batch)
+
+    assert allclose(output, reference_result)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("model_dim", [1024, 6144, 6784])
+def test_problem_size_permutations(model_dim: int) -> None:
+    """
+    Validate for different embedding sizes.
+    """
+    dtype = torch.float16
+    seq_lens = [128, 64, 192, 32]
+
+    batch = build_simple_batch(seq_lens, padding=False)
+    hidden_states = torch.randn((batch.current_tokens, model_dim),
+                                dtype=dtype,
+                                device=get_accelerator().current_device())
+
+    reference_result = baseline_implementation(hidden_states, seq_lens)
+
+    kernel = RaggedLogitsGather(model_dim, dtype)
+    output = torch.empty_like(reference_result)
+    kernel(output, hidden_states, batch)
+
+    assert allclose(output, reference_result)

tests/unit/inference/v2/kernels/ragged_ops/test_moe_gather.py

@@ -0,0 +1,122 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import DtypeEnum
+from deepspeed.inference.v2.kernels.ragged_ops import (
+    MoEGather,
+    MoEScatter,
+    RaggedTopKGating,
+)
+from .ragged_testing_utils import build_simple_batch
+"""
+For simplicity's sake, these tests do rely on ``RaggedTopKGating``  and
+``MoEScatter`` to produce correct inputs. If either of these kernels is broken
+these tests will fail, so double check the unit test results there before
+debugging here.
+"""
+
+TEST_CASES = [
+    # (n_tokens, n_experts, n_top_k)
+    (13, 64, 1),
+    (278, 64, 1),
+    (1977, 64, 1),
+    (13, 8, 2),
+    (278, 8, 2),
+    (1977, 8, 2),
+]
+
+
+def build_inputs(n_tokens: int, n_experts: int, n_top_k: int, do_padding: bool):
+
+    assert n_tokens <= 2048, "This test will break if n_tokens > 2048"
+
+    # Sequence composition shouldn't matter here
+    batch = build_simple_batch([n_tokens], padding=do_padding)
+
+    logits = torch.randn((batch.tensor_toks, n_experts),
+                         dtype=torch.float16,
+                         device=get_accelerator().current_device())
+
+    # This will make each token's value equal to its index. NOTE: This will break for
+    # tokens with index > 2048.
+    hidden_states = torch.arange(batch.tensor_toks, dtype=torch.float16,
+                                 device=get_accelerator().current_device()).repeat_interleave(4096, dim=0).reshape(
+                                     batch.tensor_toks, 4096).contiguous()
+
+    gate = RaggedTopKGating(DtypeEnum.fp16)
+
+    # Gating outputs
+    expert_counts = torch.zeros((n_experts, ), dtype=torch.int32, device=get_accelerator().current_device())
+    scores = torch.empty((batch.tensor_toks, n_top_k), dtype=torch.float32, device=get_accelerator().current_device())
+    expert_assignment = torch.empty((batch.tensor_toks, n_top_k),
+                                    dtype=torch.int32,
+                                    device=get_accelerator().current_device())
+    expert_offset = torch.empty((batch.tensor_toks, n_top_k),
+                                dtype=torch.int32,
+                                device=get_accelerator().current_device())
+
+    gate(expert_counts, scores, expert_assignment, expert_offset, logits, batch)
+
+    # Scatter outputs
+    moe_input = torch.empty((batch.tensor_toks * n_top_k, 4096),
+                            dtype=torch.float16,
+                            device=get_accelerator().current_device())
+    expert_cumsum = torch.empty((n_experts, ), dtype=torch.int64, device=get_accelerator().current_device())
+    mapped_slots = torch.empty((batch.tensor_toks, n_top_k),
+                               dtype=torch.int32,
+                               device=get_accelerator().current_device())
+
+    scatter = MoEScatter(DtypeEnum.fp16, 4096)
+    scatter(moe_input, expert_cumsum, mapped_slots, hidden_states, expert_counts, expert_assignment, expert_offset)
+
+    return batch, moe_input, scores, mapped_slots, expert_counts
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, n_experts, n_top_k", TEST_CASES)
+@pytest.mark.parametrize("do_padding", [False])
+def test_moe_gather(n_tokens: int, n_experts: int, n_top_k: int, do_padding: bool):
+    get_accelerator().manual_seed(0xC0FFEE)
+
+    batch, moe_input, scores, mapped_slots, expert_counts = build_inputs(n_tokens, n_experts, n_top_k, do_padding)
+
+    output = torch.randn((batch.tensor_toks, 4096), dtype=torch.float16, device=get_accelerator().current_device())
+
+    gather = MoEGather(DtypeEnum.fp16, 4096)
+    gather(output, moe_input, scores, mapped_slots, expert_counts)
+
+    for token_idx in range(n_tokens):
+        effective_score = scores[token_idx].sum().item()
+        assert torch.equal(
+            output[token_idx],
+            torch.full((4096, ),
+                       token_idx * effective_score,
+                       dtype=torch.float16,
+                       device=get_accelerator().current_device()))
+
+
+@pytest.mark.inference_v2_ops
+def test_moe_gather_normalize_scales():
+    get_accelerator().manual_seed(0xC0FFEE)
+
+    n_tokens = 72
+    n_experts = 8
+    n_top_k = 2
+    do_padding = False
+
+    batch, moe_input, scores, mapped_slots, expert_counts = build_inputs(n_tokens, n_experts, n_top_k, do_padding)
+    output = torch.randn((batch.tensor_toks, 4096), dtype=torch.float16, device=get_accelerator().current_device())
+
+    gather = MoEGather(DtypeEnum.fp16, 4096, normalize_scores=True)
+    gather(output, moe_input, scores, mapped_slots, expert_counts)
+
+    for token_idx in range(n_tokens):
+        assert torch.equal(
+            output[token_idx],
+            torch.full((4096, ), token_idx, dtype=torch.float16, device=get_accelerator().current_device()))

tests/unit/inference/v2/kernels/ragged_ops/test_moe_scatter.py

@@ -0,0 +1,95 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import DtypeEnum
+from deepspeed.inference.v2.kernels.ragged_ops import MoEScatter, RaggedTopKGating
+from .ragged_testing_utils import build_simple_batch
+"""
+For simplicity's sake, these tests do rely on ``RaggedTopKGating`` to produce correct
+inputs. If ``RaggedTopKGating`` is broken, these tests will fail, so double check
+the unit test results there before debugging here.
+"""
+
+TEST_CONFIGS = [
+    (13, 64, 1),
+    (278, 64, 1),
+    (1977, 64, 1),
+    (13, 8, 2),
+    (278, 8, 2),
+    (1977, 8, 2),
+]
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens, n_experts, n_top_k", TEST_CONFIGS)
+@pytest.mark.parametrize("do_padding", [False, True])
+def test_moe_scatter(n_tokens, n_experts, n_top_k, do_padding):
+
+    # Sequence composition shouldn't matter here
+    batch = build_simple_batch([n_tokens], padding=do_padding)
+
+    logits = torch.randn((batch.tensor_toks, n_experts),
+                         dtype=torch.float16,
+                         device=get_accelerator().current_device())
+
+    # This will make each token's value equal to its index. NOTE: This will break for
+    # tokens with index > 2048.
+    hidden_states = torch.arange(batch.tensor_toks, dtype=torch.float16,
+                                 device=get_accelerator().current_device()).repeat_interleave(4096, dim=0).reshape(
+                                     batch.tensor_toks, 4096).contiguous()
+
+    gate = RaggedTopKGating(DtypeEnum.fp16)
+
+    # Gating outputs
+    expert_counts = torch.zeros((n_experts, ), dtype=torch.int32, device=get_accelerator().current_device())
+    scores = torch.empty((batch.tensor_toks, n_top_k), dtype=torch.float32, device=get_accelerator().current_device())
+    expert_assignment = torch.empty((batch.tensor_toks, n_top_k),
+                                    dtype=torch.int32,
+                                    device=get_accelerator().current_device())
+    expert_offset = torch.empty((batch.tensor_toks, n_top_k),
+                                dtype=torch.int32,
+                                device=get_accelerator().current_device())
+
+    gate(expert_counts, scores, expert_assignment, expert_offset, logits, batch)
+
+    # Scatter outputs
+    moe_input = torch.empty((batch.tensor_toks * n_top_k, 4096),
+                            dtype=torch.float16,
+                            device=get_accelerator().current_device())
+    expert_cumsum = torch.empty((n_experts, ), dtype=torch.int64, device=get_accelerator().current_device())
+    mapped_slots = torch.empty((batch.tensor_toks, n_top_k),
+                               dtype=torch.int32,
+                               device=get_accelerator().current_device())
+
+    scatter = MoEScatter(DtypeEnum.fp16, 4096)
+    scatter(moe_input, expert_cumsum, mapped_slots, hidden_states, expert_counts, expert_assignment, expert_offset)
+    get_accelerator().synchronize()
+    assert torch.equal(expert_cumsum, torch.cumsum(expert_counts, dim=0).to(torch.int64))
+
+    if not do_padding:
+        assert torch.unique(mapped_slots).size(0) == n_top_k * n_tokens
+
+    for token_idx in range(batch.tensor_toks):
+        if token_idx < n_tokens:
+            for k in range(n_top_k):
+                expert_idx = expert_assignment[token_idx][k].item()
+                if expert_idx == 0:
+                    expert_cumsum_val = 0
+                else:
+                    expert_cumsum_val = expert_cumsum[expert_idx - 1]
+                offset = expert_offset[token_idx][k]
+                total_offset = offset + expert_cumsum_val
+
+                assert total_offset == mapped_slots[token_idx][k].item()
+                assert torch.equal(moe_input[total_offset], hidden_states[token_idx])
+        else:
+            for k in range(n_top_k):
+                assert mapped_slots[token_idx][k].item() == -1
+
+    assert expert_cumsum[-1] == n_tokens * n_top_k

tests/unit/inference/v2/kernels/ragged_ops/test_ragged_embed.py

@@ -0,0 +1,177 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import List, Optional, Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.kernels.ragged_ops import RaggedEmbeddingKernel
+from ....v2.inference_test_utils import allclose, get_dtypes
+from .ragged_testing_utils import build_batch_and_manager
+
+
+def baseline_implementation(token_ids: torch.Tensor,
+                            embedding_table: torch.Tensor,
+                            unpadded_size: int,
+                            positional_embedding_table: Optional[torch.Tensor] = None,
+                            positional_ids: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """
+    Baseline implementation for our ragged embedding kernel.
+    """
+    if unpadded_size == token_ids.shape[0]:
+        token_embed = torch.nn.functional.embedding(token_ids, embedding_table)
+
+        if positional_embedding_table is not None:
+            pos_embed = torch.nn.functional.embedding(positional_ids, positional_embedding_table)
+            token_embed += pos_embed
+        return token_embed
+    else:
+        real_token_ids = token_ids[:unpadded_size]
+        output = torch.empty((token_ids.shape[0], embedding_table.shape[1]),
+                             dtype=embedding_table.dtype,
+                             device=get_accelerator().current_device())
+        unpadded_output = torch.nn.functional.embedding(real_token_ids, embedding_table)
+
+        # Positional embeddings aren't padded because it's simulated
+        if positional_embedding_table is not None:
+            pos_embed = torch.nn.functional.embedding(positional_ids, positional_embedding_table)
+            unpadded_output += pos_embed
+
+        output[:unpadded_size] = unpadded_output
+        return output
+
+
+def _ragged_embed_test_helper(sequence_config: List[Tuple[int, int]],
+                              embed_dtype: torch.dtype,
+                              token_dtype: torch.dtype,
+                              embed_dim: int,
+                              vocab_size: int,
+                              do_padding: bool = False,
+                              pos_embed_size: int = -1,
+                              pos_embed_offset: int = 0) -> None:
+    """
+    Helper for embedding test to limit the number of tests to run.
+
+    Params:
+        embed_dim (int): Model dimension
+        vocab_size (int): Leading dimension on embedding weight
+        pos_embed_size (int): Size of positional embedding. If negative, no positional embedding
+            is used.
+        pos_embed_offset (int): Offset for positional embedding. Effectively, the raw offsets
+            of a token into a sequence are offset by this amount into the embedding matrix. (
+            i.e. the shape of the positional embeddings is (pos_embed_size + pos_embed_offset
+            embed_dim)
+    """
+    device = get_accelerator().current_device()
+
+    # Heads/Block size are irrelevant here but need something.
+    batch, _, _, = build_batch_and_manager(sequence_config, 64, 16, 64, vocab_range=vocab_size, padding=do_padding)
+
+    embedding_table = torch.randn((vocab_size, embed_dim), dtype=embed_dtype, device=device)
+
+    if pos_embed_size > 0:
+        pos_embedding_table = torch.randn((pos_embed_size + pos_embed_offset, embed_dim),
+                                          dtype=embed_dtype,
+                                          device=device)
+        positional_ids = torch.cat([
+            torch.arange(start_idx, start_idx + seq_len, dtype=token_dtype, device=device)
+            for seq_len, start_idx in sequence_config
+        ]) + pos_embed_offset
+    else:
+        pos_embedding_table = None
+        positional_ids = None
+
+    baseline_output = baseline_implementation(batch.input_ids().to(token_dtype), embedding_table, batch.current_tokens,
+                                              pos_embedding_table, positional_ids)
+
+    kernel = RaggedEmbeddingKernel(embed_dtype, token_dtype, embed_dim)
+    output = torch.empty_like(baseline_output)
+
+    kernel(output,
+           batch,
+           embedding_table,
+           position_embed_weight=pos_embedding_table,
+           position_embed_offset=pos_embed_offset)
+
+    if do_padding:
+        assert output.shape[0] != batch.current_tokens
+
+    assert allclose(output[:batch.current_tokens], baseline_output[:batch.current_tokens])
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('token_dtype', [torch.int32, torch.int64])
+@pytest.mark.parametrize('embed_dtype', get_dtypes())
+def test_dtype_permutations(token_dtype: torch.dtype, embed_dtype: torch.dtype) -> None:
+    """
+    Validate (on a single problem size) that the kernel support for different data types
+    is correct.
+    """
+    embed_dim = 4096
+    vocab_size = 50304
+
+    _ragged_embed_test_helper([(256, 0)], embed_dtype, token_dtype, embed_dim, vocab_size)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('vocab_size, embed_dim', [(1024, 1024), (32000, 5120), (50304, 6144)])
+def test_problem_size_permutations(vocab_size: int, embed_dim: int) -> None:
+    """
+    Validate on wider range of problem sizes.
+    """
+
+    _ragged_embed_test_helper([(256, 0)], torch.float16, torch.int32, embed_dim, vocab_size)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('seq_lens', [[128, 64, 192, 32], [57, 112, 63, 89, 1, 1, 1, 1]])
+@pytest.mark.parametrize('do_padding', [True, False])
+def test_complex_sequences(seq_lens: List[int], do_padding: bool) -> None:
+    """
+    Validate on different ragged batch construction scenarios.
+    """
+    embed_dim = 4096
+    vocab_size = 50304
+
+    _ragged_embed_test_helper([(seq_len, 0) for seq_len in seq_lens],
+                              torch.float16,
+                              torch.int32,
+                              embed_dim,
+                              vocab_size,
+                              do_padding=do_padding)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("seq_lens", [[(256, 0)], [(256, 0),
+                                                   (128, 0)], [(256, 0), (128, 0),
+                                                               (64, 0)], [(1, 877), (619, 0), (213, 372), (1, 45)]])
+def test_positional_embedding(seq_lens: List[Tuple[int, int]]) -> None:
+    """
+    Validate that positional embedding works correctly.
+    """
+    embed_dim = 4096
+    vocab_size = 50304
+
+    _ragged_embed_test_helper(seq_lens, torch.float16, torch.int32, embed_dim, vocab_size, pos_embed_size=2048)
+
+
+@pytest.mark.inference_v2_ops
+def test_positional_embedding_offset() -> None:
+    """
+    Validate that positional embedding works correctly with an offset.
+    """
+    embed_dim = 4096
+    vocab_size = 50304
+    seq_config = [(1, 877), (619, 0), (213, 372), (1, 45)]
+
+    _ragged_embed_test_helper(seq_config,
+                              torch.float16,
+                              torch.int32,
+                              embed_dim,
+                              vocab_size,
+                              pos_embed_size=2048,
+                              pos_embed_offset=2)

tests/unit/inference/v2/kernels/ragged_ops/test_top_k_gating.py

@@ -0,0 +1,171 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import DtypeEnum
+from deepspeed.inference.v2.kernels.ragged_ops import RaggedTopKGating
+from .ragged_testing_utils import build_simple_batch
+from ...inference_test_utils import allclose
+
+
+def _top_k_gating_testing_helper(n_tokens: int, n_experts: int, n_top_k: int, seed: int = 0xC0FFEE) -> None:
+
+    torch.manual_seed(seed)
+    logits = torch.randn((n_tokens, n_experts), dtype=torch.float16, device=get_accelerator().current_device())
+    batch = build_simple_batch([n_tokens], padding=False)
+    gate = RaggedTopKGating(DtypeEnum.fp16)
+
+    expert_counts = torch.zeros((n_experts, ), dtype=torch.int32, device=get_accelerator().current_device())
+    scores = torch.empty((n_tokens, n_top_k), dtype=torch.float32, device=get_accelerator().current_device())
+    expert_assignment = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+    expert_offset = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+
+    gate(expert_counts, scores, expert_assignment, expert_offset, logits, batch)
+
+    ref_weights = F.softmax(logits, dim=-1, dtype=torch.float32)
+    ref_scores, ref_indices = torch.topk(ref_weights, n_top_k, dim=-1)
+
+    assert allclose(scores, ref_scores), f"expected {ref_scores}, got {scores}"
+    assert torch.equal(expert_assignment,
+                       ref_indices.to(torch.int32)), f"expected {ref_indices}, got {expert_assignment}"
+    assert expert_counts.sum(
+    ) == n_tokens * n_top_k, f"expected {n_tokens * n_top_k} tokens, got {expert_counts.sum()}"
+
+    # Ensure that the expert offsets are unique
+    for i in range(n_experts):
+        expert_idxs = torch.where(expert_assignment == i, expert_offset, 0)
+        if expert_counts[i] > 0:
+            assert expert_idxs.unique().shape[0] == expert_counts[
+                i], f"expected {expert_counts[i]} unique offsets, got {expert_idxs.unique().shape[0]}"
+            assert expert_idxs.max(
+            ) == expert_counts[i] - 1, f"expected max offset {expert_counts[i] - 1}, got {expert_idxs.max()}"
+        else:
+            # Should have all 0's so one unique value
+            assert expert_idxs.unique().shape[0] == 1
+            assert expert_idxs.max() == 0
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('n_tokens', [1, 17, 32, 89, 433])
+def test_top_2_e_8_gating(n_tokens: int) -> None:
+    _top_k_gating_testing_helper(n_tokens=n_tokens, n_experts=8, n_top_k=2)
+
+
+def _test_single_mapping_helper(n_tokens: int,
+                                n_experts: int,
+                                assigned_expert: int,
+                                logit_fill: float = 0.0,
+                                match_fill: float = 1.0) -> None:
+
+    n_top_k = 1
+    logits = torch.full((n_tokens, n_experts),
+                        logit_fill,
+                        dtype=torch.float16,
+                        device=get_accelerator().current_device())
+
+    logits[:, assigned_expert] = match_fill
+
+    gate = RaggedTopKGating(DtypeEnum.fp16)
+
+    expert_counts = torch.zeros((n_experts, ), dtype=torch.int32, device=get_accelerator().current_device())
+    scores = torch.empty((n_tokens, n_top_k), dtype=torch.float32, device=get_accelerator().current_device())
+    expert_assignment = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+    expert_offset = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+    batch = build_simple_batch([n_tokens], padding=False)
+
+    gate(expert_counts, scores, expert_assignment, expert_offset, logits, batch)
+
+    assert expert_counts[assigned_expert] == n_tokens
+    assert torch.all(expert_assignment == assigned_expert)
+    assert torch.unique(expert_offset).shape[0] == n_tokens
+    assert allclose(scores, F.softmax(logits.float(), dim=1)[:, assigned_expert].reshape(-1, n_top_k))
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('n_tokens, n_experts', [(1, 16), (17, 16), (32, 128), (89, 128), (433, 128)])
+def test_single_mapping_gating(n_tokens: int, n_experts: int) -> None:
+    """
+    Evaluate our expert stacking behavior in complete isolation. This ensures all tokens
+    mapped to the same expert are getting unique offsets and identical scores.
+    """
+    assigned_expert = 13
+    _test_single_mapping_helper(n_tokens, n_experts, assigned_expert)
+
+
+@pytest.mark.inference_v2_ops
+def test_negative_logits():
+    """
+    Ensure that scores/values are propagated correctly when all the logits are negative. An
+    earlier implementation of the scoring would return NaN for this case.
+    """
+    _test_single_mapping_helper(128, 32, 13, logit_fill=-2.0, match_fill=-1.0)
+
+
+@pytest.mark.inference_v2_ops
+def test_determinism():
+    """
+    Ensure that ties between two logits are broken deterministically. This is essential when
+    the gating is distributed across multiple devices that need to map the same token to
+    the same expert.
+    """
+
+    n_tokens = 512
+    n_experts = 64
+    n_top_k = 1
+
+    logits = torch.zeros((n_tokens, n_experts), dtype=torch.float16, device=get_accelerator().current_device())
+    batch = build_simple_batch([n_tokens], padding=False)
+
+    logits[:, 19] = 1.0
+    logits[:, 26] = 1.0
+
+    gate = RaggedTopKGating(DtypeEnum.fp16)
+
+    for _ in range(1024):
+        expert_counts = torch.zeros((n_experts, ), dtype=torch.int32, device=get_accelerator().current_device())
+        scores = torch.empty((n_tokens, n_top_k), dtype=torch.float32, device=get_accelerator().current_device())
+        expert_assignment = torch.empty((n_tokens, n_top_k),
+                                        dtype=torch.int32,
+                                        device=get_accelerator().current_device())
+        expert_offset = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+        batch = build_simple_batch([n_tokens], padding=False)
+
+        gate(expert_counts, scores, expert_assignment, expert_offset, logits, batch)
+
+        assert expert_counts[19] == n_tokens
+        assert expert_counts[26] == 0
+        assert torch.all(expert_assignment == 19)
+        assert torch.unique(expert_offset).shape[0] == n_tokens
+        assert allclose(scores, F.softmax(logits.float(), dim=1)[:, 19].reshape(-1, 1))
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize('n_tokens, n_experts', [(1, 16), (17, 16), (32, 128), (89, 128), (433, 2)])
+def test_score_accuracy(n_tokens: int, n_experts: int) -> None:
+    """
+    Validate expert scores are correct.
+    """
+    logits = torch.randn((n_tokens, n_experts), dtype=torch.float16, device=get_accelerator().current_device())
+    batch = build_simple_batch([n_tokens], padding=False)
+    n_top_k = 1
+
+    gate = RaggedTopKGating(DtypeEnum.fp16)
+
+    expert_counts = torch.zeros((n_experts, ), dtype=torch.int32, device=get_accelerator().current_device())
+    scores = torch.empty((n_tokens, n_top_k), dtype=torch.float32, device=get_accelerator().current_device())
+    expert_assignment = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+    expert_offset = torch.empty((n_tokens, n_top_k), dtype=torch.int32, device=get_accelerator().current_device())
+
+    ref_scores = F.softmax(logits.float(), dim=1).max(dim=1).values
+    ref_scores = ref_scores.reshape(-1, 1)
+
+    gate(expert_counts, scores, expert_assignment, expert_offset, logits, batch)
+
+    assert allclose(scores, ref_scores)
+    assert expert_counts.sum() == n_tokens

tests/unit/inference/v2/model_implementations/__init__.py

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

tests/unit/inference/v2/model_implementations/parameters/__init__.py

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

tests/unit/inference/v2/model_implementations/parameters/test_contiguify.py

@@ -0,0 +1,120 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import List
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.model_implementations.flat_model_helpers import (
+    flatten_inference_model,
+    restore_inference_model,
+)
+from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
+from .utils import SimpleParam, DummyInferenceModel
+
+
+class TransformerLayerContainer(LayerContainer):
+    """
+    Stub layer container
+    """
+    PARAM_MAPPING = {
+        "param_1": "param_1.param",
+        "param_2": "param_2.param",
+    }
+
+    param_1: SimpleParam
+
+    param_2: SimpleParam
+
+
+class NonTransformerContainer(LayerContainer):
+    """
+    Stub layer container
+    """
+    PARAM_MAPPING = {
+        "param_1": "param_1.param",
+        "param_2": "param_2.param",
+        "param_3": "param_3.param",
+    }
+
+    param_1: SimpleParam
+
+    param_2: SimpleParam
+
+    param_3: SimpleParam
+
+
+@pytest.mark.inference_v2
+def test_contiguify_roundtrip():
+    """
+    Validate that contiguify round trips and reconstructions are correct.
+    """
+    model = DummyInferenceModel()
+
+    n_layers = 2
+    transformer_params = []
+    transformer_containers = []
+
+    # Create parameters and populate them into the containers
+    for i in range(n_layers):
+        transformer_containers.append(TransformerLayerContainer(model))
+        layer_params = []
+        for j in range(2):
+            layer_params.append(torch.rand(16, 16))
+            transformer_containers[i].set_dependency(f"param_{j+1}", layer_params[j])
+
+        layer_params = [p.to(get_accelerator().current_device()) for p in layer_params]
+
+        transformer_params.append(layer_params)
+        assert transformer_containers[i].is_populated == True
+
+    non_transformer_params = []
+    non_transformer_container = NonTransformerContainer(model)
+
+    for i in range(3):
+        non_transformer_params.append(torch.rand(16, 16).permute(1, 0))
+        non_transformer_container.set_dependency(f"param_{i+1}", non_transformer_params[i])
+
+    non_transformer_params = [p.to(get_accelerator().current_device()) for p in non_transformer_params]
+
+    def validate_containers(t_containers: List[LayerContainer], n_t_containers: LayerContainer,
+                            t_params: List[List[torch.Tensor]], n_t_params: List[torch.Tensor]):
+        """
+        Validate params match what is on the containers.
+        """
+        for i in range(n_layers):
+            l_c = t_containers[i]
+
+            assert l_c.is_initialized == True
+
+            assert torch.equal(l_c.param_1, t_params[i][0])
+            assert torch.equal(l_c.param_2, t_params[i][1])
+
+        assert n_t_containers.is_initialized == True
+        assert torch.equal(n_t_containers.param_1, n_t_params[0])
+        assert torch.equal(n_t_containers.param_2, n_t_params[1])
+        assert torch.equal(n_t_containers.param_3, n_t_params[2])
+        assert not n_t_containers.param_1.is_contiguous()
+        assert not n_t_containers.param_2.is_contiguous()
+        assert not n_t_containers.param_3.is_contiguous()
+
+    buffer, metadata = flatten_inference_model(transformer_containers, non_transformer_container, "NoOpPolicy")
+
+    # Validate containers before contiguify
+    validate_containers(transformer_containers, non_transformer_container, transformer_params, non_transformer_params)
+
+    # Validate restore pass
+    transformer_containers_r = []
+    for i in range(n_layers):
+        transformer_containers_r.append(TransformerLayerContainer(model))
+
+    non_transformer_container_r = NonTransformerContainer(model)
+
+    restore_inference_model(buffer, metadata, transformer_containers_r, non_transformer_container_r)
+
+    validate_containers(transformer_containers_r, non_transformer_container_r, transformer_params,
+                        non_transformer_params)

tests/unit/inference/v2/model_implementations/parameters/test_layer_inheritance.py

@@ -0,0 +1,48 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.inference.v2.inference_parameter import InferenceParameter
+from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
+
+from .utils import SimpleParam, DummyInferenceModel
+
+
+class ParentLayer(LayerContainer):
+    """
+    A layer that has a dependency on a simple parameter.
+    """
+
+    param_1: SimpleParam
+
+
+class ChildLayer(ParentLayer):
+    """
+    A layer that inherits from another layer.
+    """
+
+    param_2: SimpleParam
+
+
+@pytest.mark.inference_v2
+def test_layer_inheritance():
+    inference_model = DummyInferenceModel()
+
+    multi_param_layer = ChildLayer(inference_model)
+
+    assert multi_param_layer.n_params == 2
+    assert multi_param_layer.is_initialized is False
+
+    multi_param_layer.param_1.param = torch.ones(16, 16)
+
+    assert multi_param_layer.is_initialized is False
+
+    multi_param_layer.param_2.param = torch.full((16, 16), 2.0)
+
+    assert multi_param_layer.is_populated is True
+    assert isinstance(multi_param_layer.param_1, InferenceParameter)
+    assert isinstance(multi_param_layer.param_2, InferenceParameter)

tests/unit/inference/v2/model_implementations/parameters/test_mapping.py

@@ -0,0 +1,168 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.inference.v2.allocator import on_device
+from deepspeed.inference.v2.inference_parameter import InferenceParameter
+from deepspeed.inference.v2.model_implementations.parameter_base import ParameterBase, ParamList
+from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
+
+
+class MultiDependencyContainer(ParameterBase):
+
+    dependency_1: torch.Tensor
+
+    dependency_2: torch.Tensor
+
+    @on_device
+    def finalize(self) -> torch.Tensor:
+        param = torch.cat([self.dependency_1, self.dependency_2])
+        return InferenceParameter.initialize(param)
+
+
+class ListDependencyContainer(ParameterBase):
+
+    dependencies: ParamList("list_items")  # noqa: F821
+
+    @on_device
+    def finalize(self) -> torch.Tensor:
+        param = torch.cat(tuple(self.dependencies))
+        return InferenceParameter.initialize(param)
+
+
+class MappingLayer(LayerContainer):
+    PARAM_MAPPING = {
+        "model.val.item.d_1": "multi_depend.dependency_1",
+        "model.val.item.d_2": "multi_depend.dependency_2",
+        "model.list_vals.*.d": "list_depend.dependencies"
+    }
+
+    multi_depend: MultiDependencyContainer
+
+    list_depend: ListDependencyContainer
+
+
+class SubMappingLayer(MappingLayer):
+    PARAM_MAPPING = {
+        "model.val.item2.d_1": "multi_depend2.dependency_1",
+        "model.val.item2.d_2": "multi_depend2.dependency_2",
+    }
+
+    multi_depend2: MultiDependencyContainer
+
+
+class DoubleMappingLayer(LayerContainer):
+    PARAM_MAPPING = {
+        "model.val.item.d_1": ["multi_depend.dependency_1", "multi_depend.dependency_2"],
+    }
+
+    multi_depend: MultiDependencyContainer
+
+
+class InferenceModel:
+
+    @property
+    def list_items(self) -> int:
+        return 16
+
+
+@pytest.mark.inference_v2
+def test_mapping_syntax():
+    model = InferenceModel()
+
+    mapping_layer = MappingLayer(model)
+
+    mapping_layer.set_dependency("model.val.item.d_1", torch.ones(1))
+    mapping_layer.set_dependency("model.val.item.d_2", torch.ones(1) * 2)
+
+    assert isinstance(mapping_layer.multi_depend, torch.Tensor)
+
+    for i in range(16):
+        mapping_layer.set_dependency(f"model.list_vals.{i}.d", torch.ones(1) * i)
+        if i != 16 - 1:
+            assert mapping_layer.is_populated == False
+
+    assert isinstance(mapping_layer.list_depend, InferenceParameter)
+    assert mapping_layer.is_populated == True
+
+
+@pytest.mark.inference_v2
+def test_sub_mapping_syntax():
+    model = InferenceModel()
+
+    mapping_layer = SubMappingLayer(model)
+
+    mapping_layer.set_dependency("model.val.item.d_1", torch.ones(1))
+    mapping_layer.set_dependency("model.val.item.d_2", torch.ones(1) * 2)
+
+    assert isinstance(mapping_layer.multi_depend, InferenceParameter)
+
+    mapping_layer.set_dependency("model.val.item2.d_1", torch.ones(1))
+    mapping_layer.set_dependency("model.val.item2.d_2", torch.ones(1) * 2)
+
+    assert isinstance(mapping_layer.multi_depend2, InferenceParameter)
+
+    # We want to check into double digits to make sure that this isn't specific
+    # to single difit indexing.
+    for i in range(16):
+        mapping_layer.set_dependency(f"model.list_vals.{i}.d", torch.ones(1) * i)
+        if i != 16 - 1:
+            assert mapping_layer.is_populated == False
+
+    assert isinstance(mapping_layer.list_depend, InferenceParameter)
+    assert mapping_layer.is_populated == True
+
+
+@pytest.mark.inference_v2
+def test_double_mapping_syntax():
+    model = InferenceModel()
+
+    mapping_layer = DoubleMappingLayer(model)
+    mapping_layer.set_dependency("model.val.item.d_1", torch.ones(1))
+
+    # The single parameter setting should immediately make the parameter finalized
+    # and the whole layer initialized.
+    assert isinstance(mapping_layer.multi_depend, InferenceParameter)
+    assert mapping_layer.is_populated == True
+
+
+@pytest.mark.inference_v2
+def test_insufficient_mapping_syntax():
+    """
+    In the above example, we don't have a mapping for `multi_depend2.dependency_2`.
+    """
+
+    with pytest.raises(ValueError):
+
+        class InsuffienctMappingLayer(LayerContainer):
+            PARAM_MAPPING = {
+                "model.val.item.d_1": "multi_depend1.dependency_1",
+                "model.val.item.d_2": "multi_depend1.dependency_2",
+                "model.val.item2.d_1": "multi_depend2.dependency_1",
+            }
+
+            multi_depend1: MultiDependencyContainer
+
+            multi_depend2: MultiDependencyContainer
+
+
+@pytest.mark.inference_v2
+def test_unknown_target_mapping_syntax():
+    """
+    In the above example, `multi_depend_unknown` does not exist
+    """
+
+    with pytest.raises(ValueError):
+
+        class UnknownTargetMappingLayer(LayerContainer):
+            PARAM_MAPPING = {
+                "model.val.item.d_1": "multi_depend1.dependency_1",
+                "model.val.item.d_2": "multi_depend1.dependency_2",
+                "model.val.item2.d_1": "multi_depend_unknown.dependency_1",
+            }
+
+            multi_depend: MultiDependencyContainer

tests/unit/inference/v2/model_implementations/parameters/test_multi_parameter_layer.py

@@ -0,0 +1,79 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.inference.v2.inference_parameter import InferenceParameter
+from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
+
+from .utils import validate_device, SimpleParam, ListParam, DummyInferenceModel
+
+
+class MultiParameterLayer(LayerContainer):
+    """
+    Two dependencies, both of which are simple parameters.
+    """
+
+    param_1: SimpleParam
+
+    param_2: SimpleParam
+
+
+class MixedMultiParameterLayer(LayerContainer):
+    """
+    Two dependencies, one of which is a simple parameter, the other is a list parameter.
+    """
+
+    param_1: SimpleParam
+
+    param_2: ListParam
+
+
+@pytest.mark.inference_v2
+def test_multi_parameter_layer():
+    inference_model = DummyInferenceModel()
+
+    multi_param_layer = MultiParameterLayer(inference_model)
+
+    assert multi_param_layer.n_params == 2
+    assert multi_param_layer.is_populated is False
+
+    multi_param_layer.param_1.param = torch.ones(16, 16)
+
+    assert multi_param_layer.is_populated is False
+
+    multi_param_layer.param_2.param = torch.full((16, 16), 2.0)
+
+    assert multi_param_layer.is_populated is True
+    assert isinstance(multi_param_layer.param_1, InferenceParameter)
+    assert isinstance(multi_param_layer.param_2, InferenceParameter)
+
+
+@pytest.mark.inference_v2
+def test_mixed_multi_parameter_layer():
+    inference_model = DummyInferenceModel()
+
+    mixed_multi_param_layer = MixedMultiParameterLayer(inference_model)
+
+    assert mixed_multi_param_layer.n_params == 2
+    assert mixed_multi_param_layer.is_populated is False
+
+    mixed_multi_param_layer.param_2.params[1] = torch.full((16, 16), 2.0)
+    assert mixed_multi_param_layer.is_populated is False
+    assert not isinstance(mixed_multi_param_layer.param_2, InferenceParameter)
+
+    mixed_multi_param_layer.param_1.param = torch.ones(16, 16)
+    assert mixed_multi_param_layer.is_populated is False
+    assert isinstance(mixed_multi_param_layer.param_1, InferenceParameter)
+
+    validate_device(mixed_multi_param_layer.param_1)
+
+    mixed_multi_param_layer.param_2.params[0] = torch.full((16, 16), 2.0)
+
+    assert mixed_multi_param_layer.is_populated is True
+    assert isinstance(mixed_multi_param_layer.param_2, InferenceParameter)
+
+    validate_device(mixed_multi_param_layer.param_2)

tests/unit/inference/v2/model_implementations/parameters/test_parameter_list.py

@@ -0,0 +1,105 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.inference.v2.allocator import on_device
+from deepspeed.inference.v2.inference_parameter import InferenceParameter
+from deepspeed.inference.v2.model_implementations.parameter_base import ParameterBase, ParamList
+from deepspeed.inference.v2.model_implementations.layer_container_base import LayerContainer
+from deepspeed.inference.v2.model_implementations.common_parameters import *
+
+from .utils import validate_device
+
+
+class SimpleMoELayer(LayerContainer):
+
+    moe_mlp_1: UnfusedMoEMLP1Parameter
+
+
+class DummyInferenceModel:
+
+    def __init__(self, experts_per_rank: int) -> None:
+        self._num_experts = experts_per_rank
+
+    @property
+    def n_experts(self) -> int:
+        return self._num_experts
+
+    @on_device
+    def transform_moe_mlp_1_param(self, param: torch.Tensor) -> torch.Tensor:
+        return InferenceParameter.initialize(param)
+
+
+@pytest.mark.inference_v2
+def test_simple_moe_layer():
+
+    inference_model = DummyInferenceModel(experts_per_rank=2)
+
+    simple_moe_layer = SimpleMoELayer(inference_model)
+
+    assert simple_moe_layer.moe_mlp_1.experts[0] is None
+    assert simple_moe_layer.moe_mlp_1.experts[1] is None
+
+    # Set the first expert
+    simple_moe_layer.moe_mlp_1.experts[0] = torch.zeros(16, 16)
+
+    assert simple_moe_layer.moe_mlp_1.experts[0] is not None
+    assert simple_moe_layer.moe_mlp_1.experts[1] is None
+
+    assert not simple_moe_layer.is_initialized
+
+    # Set the second expert
+    simple_moe_layer.moe_mlp_1.experts[1] = torch.ones(16, 16)
+
+    # We have all the experts, so the layer should be initialized
+    assert simple_moe_layer.is_initialized
+    assert isinstance(simple_moe_layer.moe_mlp_1, torch.Tensor)
+
+    validate_device(simple_moe_layer.moe_mlp_1)
+
+
+"""
+Check that we can mix the number of elements in lists in the same context and have that
+be tracked correctly.
+"""
+
+
+class CustomListParam1(ParameterBase):
+
+    deps: ParamList("attr_1")
+
+
+class CustomListParam2(ParameterBase):
+
+    deps: ParamList("attr_2")
+
+
+class MixedLayer(LayerContainer):
+
+    list_1: CustomListParam1
+    list_2: CustomListParam2
+
+
+class MixedInferenceModel:
+
+    @property
+    def attr_1(self) -> int:
+        return 1
+
+    @property
+    def attr_2(self) -> int:
+        return 2
+
+
+@pytest.mark.inference_v2
+def test_mixed_param_lists():
+    model = MixedInferenceModel()
+
+    layer = MixedLayer(model)
+
+    assert layer.list_1.deps.n_params == 1
+    assert layer.list_2.deps.n_params == 2

tests/unit/inference/v2/model_implementations/parameters/utils.py

@@ -0,0 +1,60 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.allocator import on_device
+from deepspeed.inference.v2.inference_parameter import InferenceParameter
+from deepspeed.inference.v2.model_implementations.parameter_base import ParameterBase, ParametrizedList
+
+
+class SimpleParam(ParameterBase):
+    """
+    Parameter with single dependency.
+    """
+
+    param: torch.Tensor
+
+    @on_device
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform(self.param)
+
+
+class SimpleParametrizedList(ParametrizedList):
+    """
+    Parameter list based on `num_dependencies` attribute.
+    """
+
+    count_attr: str = "num_dependencies"
+
+
+class ListParam(ParameterBase):
+    """
+    Parameter with list dependency.
+
+    NOTE: This uses the tuple workaround for the `ParametrizedList` class
+    as described in the docstring of `ParametrizedList`.
+    """
+
+    params: SimpleParametrizedList
+
+    @on_device
+    def finalize(self) -> torch.Tensor:
+        return self.inference_model.transform(torch.cat(tuple(self.params)))
+
+
+class DummyInferenceModel:
+
+    @property
+    def num_dependencies(self) -> int:
+        return 2
+
+    def transform(self, param: torch.Tensor) -> torch.Tensor:
+        return InferenceParameter.initialize(param)
+
+
+def validate_device(tensor: torch.Tensor):
+    assert tensor.device == torch.device(get_accelerator().current_device())

tests/unit/inference/v2/model_implementations/sharding/__init__.py

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

tests/unit/inference/v2/model_implementations/sharding/test_attn_out_sharding.py

@@ -0,0 +1,129 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.model_implementations.sharding import *
+
+# None of the logic should be dependent on head size.
+HEAD_SIZE = 64
+
+
+def fill_with_head_ids(head_size: int, n_heads: int) -> torch.Tensor:
+    """
+    Fills a tensor with the associated head ids. All columns should have the same value.
+    """
+    head_ids = torch.arange(n_heads, dtype=torch.half, device=get_accelerator().current_device())
+
+    head_ids = head_ids.repeat_interleave(head_size).repeat(head_size * n_heads).reshape(n_heads * head_size, -1)
+    return head_ids
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("n_heads, n_shards", [(1, 1), (8, 4), (32, 8)])
+def test_mha_even_sharding(n_heads: int, n_shards: int):
+    """
+    Even head sharding for MHA.
+
+    Args:
+        n_heads (int): The number QKV heads.
+        n_shards (int): The number of shards to test for.
+    """
+    param = fill_with_head_ids(HEAD_SIZE, n_heads)
+
+    n_local_heads = n_heads // n_shards
+    sharded_shape = (HEAD_SIZE * n_heads, HEAD_SIZE * n_local_heads)
+
+    for shard_rank in range(n_shards):
+        sharded_param = shard_attn_out_param(param, shard_rank, n_shards, HEAD_SIZE)
+        n_heads_local_q, _ = get_local_heads(shard_rank, n_shards, n_heads)
+
+        assert sharded_param.shape[-1] == HEAD_SIZE * n_heads_local_q
+        assert sharded_param.shape == sharded_shape
+
+        heads = torch.chunk(sharded_param, n_local_heads, dim=1)
+
+        for i, head in enumerate(heads):
+            assert torch.all(head == i + shard_rank * n_local_heads)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("n_heads, n_shards", [(3, 2), (20, 8)])
+def test_mha_unbalanced_sharding(n_heads: int, n_shards: int):
+    """
+    Unbalanced head sharding for MHA.
+
+    Args:
+        n_heads (int): The number QKV heads.
+        n_shards (int): The number of shards to test for.
+    """
+    param = fill_with_head_ids(HEAD_SIZE, n_heads)
+
+    max_heads = 0
+    min_heads = n_heads
+    seen_heads = set()
+    total_heads = 0
+
+    for shard_rank in range(n_shards):
+        sharded_param = shard_attn_out_param(param, shard_rank, n_shards, HEAD_SIZE)
+        n_heads_local_q, _ = get_local_heads(shard_rank, n_shards, n_heads)
+
+        assert sharded_param.shape[-1] == HEAD_SIZE * n_heads_local_q
+
+        n_local_heads = sharded_param.shape[1] // HEAD_SIZE
+        total_heads += n_local_heads
+        max_heads = max(max_heads, n_local_heads)
+        min_heads = min(min_heads, n_local_heads)
+
+        for i in range(n_local_heads):
+            head_ids = torch.unique_consecutive(sharded_param[:, i * HEAD_SIZE:(i + 1) * HEAD_SIZE])
+            assert len(head_ids) == 1
+            seen_heads.add(head_ids.item())
+
+    assert max_heads == min_heads + 1
+    assert total_heads == n_heads
+    assert len(seen_heads) == n_heads
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("n_heads_q, n_heads_kv, n_shards", [(20, 4, 8)])
+def test_gqa_uneven_sharding(n_heads_q: int, n_heads_kv: int, n_shards: int):
+    """
+    We only test the uneven GQA test case because even GQA shards the attention output
+    in the exact same manner as MHA.
+
+    Args:
+        n_heads_q (int): The number of query heads.
+        n_heads_kv (int): The number of key/value heads.
+        n_shards (int): The number of shards to test for.
+    """
+    param = fill_with_head_ids(HEAD_SIZE, n_heads_q)
+
+    min_heads = n_heads_q
+    max_heads = 0
+    seen_heads = set()
+    total_heads = 0
+
+    for shard_rank in range(n_shards):
+        sharded_param = shard_attn_out_param(param, shard_rank, n_shards, HEAD_SIZE, n_heads_q, n_heads_kv)
+        n_heads_local_q, _ = get_local_heads(shard_rank, n_shards, n_heads_q, n_heads_kv)
+
+        assert sharded_param.shape[-1] == HEAD_SIZE * n_heads_local_q
+
+        n_local_heads = sharded_param.shape[1] // HEAD_SIZE
+        total_heads += n_local_heads
+        max_heads = max(max_heads, n_local_heads)
+        min_heads = min(min_heads, n_local_heads)
+
+        for i in range(n_local_heads):
+            head_id = torch.unique_consecutive(sharded_param[:, i * HEAD_SIZE:(i + 1) * HEAD_SIZE])
+            assert len(head_id) == 1
+            seen_heads.add(head_id.item())
+
+    assert max_heads == min_heads + 1
+    assert total_heads == n_heads_q
+    assert len(seen_heads) == n_heads_q

tests/unit/inference/v2/model_implementations/sharding/test_mlp_sharding.py

@@ -0,0 +1,116 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.model_implementations.sharding import *
+
+
+def round_up_to_256(x: int) -> int:
+    """
+    Round up to the nearest multiple of 256.
+    """
+    return x + (256 - x % 256)
+
+
+def make_params(model_dim: int, ffn_multiplier: int, n_experts: int, gated: bool = False) -> torch.Tensor:
+    """
+
+    """
+    if gated:
+        mlp_1_intermediate = round_up_to_256(int(model_dim * ffn_multiplier * 4 / 3))
+        mlp_2_intermediate = mlp_1_intermediate // 2
+    else:
+        mlp_1_intermediate = ffn_multiplier * model_dim
+        mlp_2_intermediate = ffn_multiplier * model_dim
+
+    mlp_1_shared_dim = torch.arange(mlp_1_intermediate, dtype=torch.float32, device=get_accelerator().current_device())
+
+    mlp_1_w = mlp_1_shared_dim.repeat_interleave(model_dim).reshape(mlp_1_intermediate, model_dim)
+    mlp_1_b = mlp_1_shared_dim
+
+    mlp_2_shared_dim = torch.arange(mlp_2_intermediate, dtype=torch.float32, device=get_accelerator().current_device())
+    mlp_2_w = mlp_2_shared_dim.repeat(model_dim).reshape(model_dim, mlp_2_intermediate)
+    mlp_2_b = torch.ones(model_dim, dtype=torch.float32, device=get_accelerator().current_device())
+
+    if n_experts > 1:
+        mlp_1_w = mlp_1_w.expand(n_experts, -1, -1)
+        mlp_1_b = mlp_1_b.expand(n_experts, -1)
+        mlp_2_w = mlp_2_w.expand(n_experts, -1, -1)
+        mlp_2_b = mlp_2_b.expand(n_experts, -1)
+
+    return (mlp_1_w, mlp_1_b, mlp_2_w, mlp_2_b)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("model_dim, ffn_multiplier, n_shards", [(1024, 4, 1), (1024, 4, 8), (1024, 4, 6)])
+@pytest.mark.parametrize("n_experts", [1, 16])
+def test_even_ffn_sharding(model_dim: int, ffn_multiplier: int, n_shards: int, n_experts: int):
+    """
+    FFN sharding tends to be much simpler than attention sharding since it works on larger granularities.
+    While the test case of (1024, 4, 6) is not a use case we're likely to see, this does ensure that
+    the sharding logic will round correctly for the alignments we care about.
+    """
+    mlp_1_w, mlp_1_b, mlp_2_w, mlp_2_b = make_params(model_dim, ffn_multiplier, n_experts)
+
+    total_ffn_dim = model_dim * ffn_multiplier
+    mapped_neurons = 0
+
+    is_moe = n_experts > 1
+
+    for shard_rank in range(n_shards):
+        shard_1_w = shard_mlp_1_param(mlp_1_w, shard_rank, n_shards, is_moe=is_moe)
+        shard_1_b = shard_mlp_1_param(mlp_1_b, shard_rank, n_shards, is_moe=is_moe)
+        shard_2_w = shard_mlp_2_param(mlp_2_w, shard_rank, n_shards, is_moe=is_moe)
+        shard_2_b = shard_mlp_2_param(mlp_2_b, shard_rank, n_shards, is_moe=is_moe)
+
+        assert shard_1_w.shape[-2] == shard_2_w.shape[-1]
+        assert shard_1_w.shape[-2] % DEFAULT_SHARD_GRANULARITY == 0
+        assert shard_1_w.shape[-2] == shard_1_b.shape[-1]
+
+        mapped_neurons += shard_1_w.shape[-2]
+
+        if shard_rank != 0:
+            assert shard_2_b is None
+        else:
+            assert shard_2_b.shape[-1] == model_dim
+
+    assert mapped_neurons == total_ffn_dim
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("model_dim, ffn_multiplier, n_shards", [(1024, 4, 1), (1024, 4, 8), (1024, 4, 6)])
+@pytest.mark.parametrize("n_experts", [1, 16])
+def test_gated_ffn_sharding(model_dim: int, ffn_multiplier: int, n_shards: int, n_experts: int):
+    """
+    Test the same cases assuming a gated regime.
+    """
+    mlp_1_w, mlp_1_b, mlp_2_w, mlp_2_b = make_params(model_dim, ffn_multiplier, n_experts, gated=True)
+
+    total_ffn_dim = round_up_to_256(int(model_dim * ffn_multiplier * 4 / 3))
+    mapped_neurons = 0
+
+    is_moe = n_experts > 1
+
+    for shard_rank in range(n_shards):
+        shard_1_w = shard_mlp_1_param(mlp_1_w, shard_rank, n_shards, gated=True, is_moe=is_moe)
+        shard_1_b = shard_mlp_1_param(mlp_1_b, shard_rank, n_shards, gated=True, is_moe=is_moe)
+        shard_2_w = shard_mlp_2_param(mlp_2_w, shard_rank, n_shards, is_moe=is_moe)
+        shard_2_b = shard_mlp_2_param(mlp_2_b, shard_rank, n_shards, is_moe=is_moe)
+
+        assert shard_1_w.shape[-2] == shard_2_w.shape[-1] * 2
+        assert shard_1_w.shape[-2] % DEFAULT_SHARD_GRANULARITY == 0
+        assert shard_1_w.shape[-2] == shard_1_b.shape[-1]
+
+        mapped_neurons += shard_1_w.shape[-2]
+
+        if shard_rank != 0:
+            assert shard_2_b is None
+        else:
+            assert shard_2_b.shape[-1] == model_dim
+
+    assert mapped_neurons == total_ffn_dim

tests/unit/inference/v2/model_implementations/sharding/test_qkv_sharding.py

@@ -0,0 +1,251 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.model_implementations.sharding import *
+
+
+def fill_with_head_ids(head_size: int, n_heads_q: int, n_heads_kv: Optional[int] = None) -> torch.Tensor:
+    """
+
+    """
+    head_ids_q = torch.arange(n_heads_q, dtype=torch.half, device=get_accelerator().current_device())
+    head_vals_q = head_ids_q.repeat_interleave(head_size * head_size * n_heads_q).reshape(n_heads_q * head_size, -1)
+
+    if n_heads_kv is None:
+        return torch.cat([head_vals_q, head_vals_q, head_vals_q], dim=0)
+
+    head_ids_k = torch.arange(n_heads_kv, dtype=torch.half, device=get_accelerator().current_device())
+    head_vals_k = head_ids_k.repeat_interleave(head_size * head_size * n_heads_q).reshape(n_heads_kv * head_size, -1)
+
+    return torch.cat([head_vals_q, head_vals_k, head_vals_k], dim=0)
+
+
+def validate_inferred_shape(shard: torch.Tensor, head_size: int, n_local_q_heads: int, n_local_kv_heads: int):
+    """
+    Validate that the leading dim of the shard is of the expected size and aligns with the sharding
+    logic for the attention computation itself.
+    """
+    inferred_leading_dim = head_size * (n_local_q_heads + 2 * n_local_kv_heads)
+    assert shard.shape[0] == inferred_leading_dim
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("n_heads,n_shards", [(1, 1), (32, 1), (32, 8)])
+def test_even_mha_sharding(head_size: int, n_heads: int, n_shards: int):
+    """
+    Test for MHA sharding. In these scenarios, we expect that each of the shards
+    should be the same size.
+    """
+    param = fill_with_head_ids(head_size, n_heads)
+
+    heads_per_shard = n_heads // n_shards
+
+    for shard_rank in range(n_shards):
+
+        shard = shard_qkv_param(param, shard_rank, n_shards, head_size, n_heads, n_heads)
+        n_local_q_heads, n_local_kv_heads = get_local_heads(shard_rank, n_shards, n_heads, n_heads)
+        validate_inferred_shape(shard, head_size, n_local_q_heads, n_local_kv_heads)
+
+        assert shard.shape == (3 * head_size * heads_per_shard, head_size * n_heads)
+
+        heads = shard.chunk(heads_per_shard * 3, dim=0)
+        for i in range(heads_per_shard):
+            assert torch.all(heads[i] == i + shard_rank * heads_per_shard)
+            assert torch.all(heads[i + heads_per_shard] == i + shard_rank * heads_per_shard)
+            assert torch.all(heads[i + heads_per_shard * 2] == i + shard_rank * heads_per_shard)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("n_heads, n_shards", [(3, 2), (20, 8)])
+def test_unbalanced_mha_sharding(head_size: int, n_heads: int, n_shards: int):
+    """
+    Test MHA sharding when the distribution of heads will not be equal across all ranks.
+    """
+    param = fill_with_head_ids(head_size, n_heads)
+
+    max_heads = 0
+    min_heads = n_heads
+    total_heads = 0
+    seen_heads = set()
+
+    for shard_rank in range(n_shards):
+        shard = shard_qkv_param(param, shard_rank, n_shards, head_size, n_heads, n_heads)
+        n_local_q_heads, n_local_kv_heads = get_local_heads(shard_rank, n_shards, n_heads, n_heads)
+        validate_inferred_shape(shard, head_size, n_local_q_heads, n_local_kv_heads)
+
+        n_heads_in_shard = shard.shape[0] // head_size // 3
+
+        max_heads = max(max_heads, n_heads_in_shard)
+        min_heads = min(min_heads, n_heads_in_shard)
+        total_heads += n_heads_in_shard
+
+        heads = shard.chunk(n_heads_in_shard * 3, dim=0)
+
+        for local_head_id in range(n_heads_in_shard):
+            head_qkv = torch.cat([
+                heads[local_head_id], heads[local_head_id + n_heads_in_shard],
+                heads[local_head_id + 2 * n_heads_in_shard]
+            ],
+                                 dim=0)
+            assert head_qkv.shape == (3 * head_size, head_size * n_heads)
+
+            global_head_id = torch.unique_consecutive(head_qkv)
+            assert len(global_head_id) == 1
+
+            seen_heads.add(global_head_id.item())
+
+    assert max_heads - min_heads <= 1
+    assert total_heads == n_heads
+    assert len(seen_heads) == n_heads
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("n_heads_q, n_heads_kv, n_shards", [(4, 2, 1), (8, 2, 1), (64, 16, 8)])
+def test_gqa_even_sharding(head_size: int, n_heads_q: int, n_heads_kv: int, n_shards: int):
+    """
+    Test GQA sharding when the KV heads are evenly divisible by the number of shards.
+    """
+    param = fill_with_head_ids(head_size, n_heads_q, n_heads_kv)
+
+    n_kv_heads_in_shard = n_heads_kv // n_shards
+    n_q_heads_in_shard = n_heads_q // n_shards
+
+    for shard_rank in range(n_shards):
+        shard = shard_qkv_param(param, shard_rank, n_shards, head_size, n_heads_q, n_heads_kv)
+        n_local_q_heads, n_local_kv_heads = get_local_heads(shard_rank, n_shards, n_heads_q, n_heads_kv)
+        validate_inferred_shape(shard, head_size, n_local_q_heads, n_local_kv_heads)
+
+        assert shard.shape[0] == (n_q_heads_in_shard + n_kv_heads_in_shard * 2) * head_size
+
+        q = shard[:n_q_heads_in_shard * head_size]
+        k = shard[n_q_heads_in_shard * head_size:(n_q_heads_in_shard + n_kv_heads_in_shard) * head_size]
+        v = shard[(n_q_heads_in_shard + n_kv_heads_in_shard) * head_size:]
+
+        for local_head_id in range(n_q_heads_in_shard):
+            assert torch.all(q[local_head_id * head_size:(local_head_id + 1) * head_size] == local_head_id +
+                             shard_rank * n_q_heads_in_shard)
+
+        for local_head_id in range(n_kv_heads_in_shard):
+            assert torch.all(k[local_head_id * head_size:(local_head_id + 1) * head_size] == local_head_id +
+                             shard_rank * n_kv_heads_in_shard)
+            assert torch.all(v[local_head_id * head_size:(local_head_id + 1) * head_size] == local_head_id +
+                             shard_rank * n_kv_heads_in_shard)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("n_heads_q, n_heads_kv, n_shards", [(4, 2, 4), (20, 4, 8)])
+def test_gqa_uneven_sharding(head_size: int, n_heads_q: int, n_heads_kv: int, n_shards: int):
+    """
+    Test GQA sharding when there are more shards than KV heads.
+    """
+    param = fill_with_head_ids(head_size, n_heads_q, n_heads_kv)
+
+    n_kv_heads_in_shard = 1
+    n_shards_per_kv_head = n_shards // n_heads_kv
+
+    max_heads = 0
+    min_heads = n_heads_q
+    total_heads = 0
+    seen_heads = set()
+
+    for shard_rank in range(n_shards):
+        shard = shard_qkv_param(param, shard_rank, n_shards, head_size, n_heads_q, n_heads_kv)
+        n_local_q_heads, n_local_kv_heads = get_local_heads(shard_rank, n_shards, n_heads_q, n_heads_kv)
+        validate_inferred_shape(shard, head_size, n_local_q_heads, n_local_kv_heads)
+
+        local_n_heads_q = (shard.shape[0] - 2 * n_kv_heads_in_shard * head_size) // head_size
+
+        max_heads = max(max_heads, local_n_heads_q)
+        min_heads = min(min_heads, local_n_heads_q)
+        total_heads += local_n_heads_q
+
+        q = shard[:local_n_heads_q * head_size]
+        kv = shard[local_n_heads_q * head_size:]
+
+        for local_head_id in range(local_n_heads_q):
+            q_head_id = torch.unique_consecutive(q[local_head_id * head_size:(local_head_id + 1) * head_size])
+            assert len(q_head_id) == 1
+
+            seen_heads.add(q_head_id.item())
+
+        kv_id_calc = shard_rank // n_shards_per_kv_head
+        kv_id = torch.unique_consecutive(kv)
+        assert len(kv_id) == 1
+        assert kv_id.item() == kv_id_calc
+
+    assert max_heads - min_heads <= 1
+    assert total_heads == n_heads_q
+    assert len(seen_heads) == n_heads_q
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("n_heads, n_shards", [(6, 8)])
+def test_unsupported_mha_configs(head_size: int, n_heads: int, n_shards: int):
+    """
+    Sharding should fail if there are fewer heads than shards.
+
+    TODO(cmikeh2): Look to support this configuration.
+    """
+    param = fill_with_head_ids(head_size, n_heads)
+
+    for shard_rank in range(n_shards):
+        with pytest.raises(ValueError):
+            shard_qkv_param(param, shard_rank, n_shards, head_size, n_heads, n_heads)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize("head_size", [64])
+@pytest.mark.parametrize("n_heads_q, n_heads_kv, n_shards", [(5, 2, 1), (40, 10, 8), (30, 5, 8)])
+def test_unsupported_gqa_configs(head_size: int, n_heads_q: int, n_heads_kv: int, n_shards: int):
+    """
+    GQA has stricter requirements. We must be able to evenly shard or distribute the KV heads.
+
+    Test cases are to test the following preconditions specifically:
+        1. n_heads_q % n_heads_kv == 0
+        2. We must be able to evenly distribute KV heads
+        3. We must be able to evely split KV heads
+    """
+    param = fill_with_head_ids(head_size, n_heads_q, n_heads_kv)
+
+    for shard_rank in range(n_shards):
+        with pytest.raises(ValueError):
+            shard_qkv_param(param, shard_rank, n_shards, head_size, n_heads_q, n_heads_kv)
+
+
+@pytest.mark.inference_v2
+def test_mha_input_shape_error():
+
+    param = torch.empty(256, 128)
+
+    n_heads = 2
+    head_size = 64
+
+    with pytest.raises(ValueError):
+        shard_qkv_param(param, 0, 1, 64)
+
+
+@pytest.mark.inference_v2
+def test_gqa_input_shape_error():
+
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 4
+
+    # Correct shape is 1536 (=16 * 64 + 2 * 4 * 64), 1024
+    param = torch.empty(2048, 1024)
+
+    with pytest.raises(ValueError):
+        shard_qkv_param(param, 0, 1, head_size, n_heads_q, n_heads_kv)

tests/unit/inference/v2/modules/__init__.py

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

tests/unit/inference/v2/modules/test_blas_linear_module.py

@@ -0,0 +1,111 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum, is_gated
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSLinearConfig
+from deepspeed.inference.v2.modules.interfaces import DSLinearRegistry
+from ...v2.inference_test_utils import allclose
+
+
+def reference_implementation(hidden_states: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor],
+                             act_type: ActivationType) -> torch.Tensor:
+    dtype = hidden_states.dtype
+    out_states = torch.nn.functional.linear(hidden_states, weight, bias)
+    out_states.float()
+
+    if is_gated(act_type):
+        act_func_map = {
+            ActivationType.ReGLU: torch.nn.functional.relu,
+            ActivationType.GEGLU: lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
+            ActivationType.SiGLU: torch.nn.functional.silu,
+        }
+
+        act_act = out_states[..., ::2]
+        act_linear = out_states[..., 1::2]
+
+        act_act = act_func_map[act_type](act_act)
+        out_states = act_act * act_linear
+    else:
+        act_func_map = {
+            ActivationType.RELU: torch.nn.functional.relu,
+            ActivationType.GELU: torch.nn.functional.gelu,
+            ActivationType.SILU: torch.nn.functional.silu,
+            ActivationType.IDENTITY: lambda x: x,
+        }
+
+        out_states = act_func_map[act_type](out_states)
+    return out_states.to(dtype)
+
+
+def _blas_linear_helper(tokens: int,
+                        in_channels: int,
+                        out_channels: int,
+                        dtype: DtypeEnum,
+                        act_fn: ActivationType,
+                        use_bias: bool = True) -> None:
+    linear_config = DSLinearConfig(max_tokens=2048,
+                                   in_channels=in_channels,
+                                   out_channels=out_channels,
+                                   activation=act_fn,
+                                   input_dtype=dtype,
+                                   output_dtype=dtype)
+
+    bundle = ConfigBundle(name='blas_fp_linear', config=linear_config)
+
+    module = DSLinearRegistry.instantiate_config(bundle)
+
+    # Input vals
+    hidden_states = torch.randn(
+        (tokens, in_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
+
+    weight_out_channels = 2 * out_channels if is_gated(act_fn) else out_channels
+    weight = torch.randn(
+        (weight_out_channels, in_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
+    if use_bias:
+        bias = torch.randn(
+            (weight_out_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
+    else:
+        bias = None
+
+    # Reference output
+    ref_output = reference_implementation(hidden_states, weight, bias, act_fn)
+
+    # New output
+    ds_output = module(hidden_states, weight, bias)
+
+    # Check
+    assert allclose(ds_output, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, in_channels, out_channels", [(1, 4608, 1728), (37, 8192, 4096), (1280, 3072, 6144)])
+def test_blas_linear_shapes(tokens: int, in_channels: int, out_channels: int) -> None:
+
+    _blas_linear_helper(tokens, in_channels, out_channels, DtypeEnum.fp16, ActivationType.IDENTITY)
+
+
+all_acts = [
+    ActivationType.RELU,
+    ActivationType.GELU,
+    ActivationType.SILU,
+    ActivationType.GEGLU,
+    ActivationType.ReGLU,
+    ActivationType.SiGLU,
+]
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("act_fn", all_acts)
+@pytest.mark.parametrize("use_bias", [True, False])
+def test_blas_linear_act_fn(act_fn: ActivationType, use_bias: bool) -> None:
+
+    _blas_linear_helper(283, 512, 4096, DtypeEnum.fp16, act_fn, use_bias=use_bias)

tests/unit/inference/v2/modules/test_blocked_attn.py

@@ -0,0 +1,207 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import itertools
+
+from typing import List, Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSSelfAttentionConfig, PositionalEmbeddingType, RotateHalfConfig
+from deepspeed.inference.v2.modules.interfaces import DSSelfAttentionRegistry, DSSelfAttentionBase
+
+from ..kernels.ragged_ops.ragged_testing_utils import build_batch_and_manager
+from ...v2.inference_test_utils import allclose
+
+try:
+    from flash_attn.flash_attn_interface import flash_attn_varlen_func
+    validate_accuracy = True
+except ImportError:
+    validate_accuracy = False
+
+
+def _blocked_flash_testing_helper(head_size: int,
+                                  n_heads_q: int,
+                                  n_heads_kv: int,
+                                  seq_params: List[Tuple[int, int]],
+                                  trained_freqs: bool = None) -> None:
+    """
+    Helper function for testing blocked flash attention. This implementation is based on
+    the implemnentation in ``unit.inference.kernels.ragged_ops.test_blocked_flash`` but
+    integrates functionality to validate the composability.
+    """
+    if trained_freqs is None:
+        embed_type = PositionalEmbeddingType.none
+        embed_args = None
+    else:
+        embed_type = PositionalEmbeddingType.rotate_half
+        embed_args = RotateHalfConfig(use_trained_freqs=trained_freqs)
+
+    attn_config = DSSelfAttentionConfig(max_tokens=2048,
+                                        n_heads_q=n_heads_q,
+                                        n_heads_kv=n_heads_kv,
+                                        head_size=head_size,
+                                        max_sequences=32,
+                                        positional_embedding_type=embed_type,
+                                        positional_embedding_config=embed_args)
+
+    config = ConfigBundle(name='dense_blocked_attention', config=attn_config)
+    attn_module: DSSelfAttentionBase = DSSelfAttentionRegistry.instantiate_config(config)
+
+    kv_block_size = attn_module.kv_block_size
+
+    kvs = []
+    for _, history_len in seq_params:
+        if history_len > 0:
+            kvs.append(
+                torch.randn((history_len, 2 * n_heads_kv * head_size),
+                            device=get_accelerator().current_device(),
+                            dtype=torch.float16))
+        else:
+            kvs.append(None)
+
+    batch, state_manager, _ = build_batch_and_manager(seq_params, head_size, n_heads_kv, kv_block_size, kv_fill=kvs)
+
+    qkv = torch.randn((batch.current_tokens, (n_heads_q + 2 * n_heads_kv) * head_size),
+                      device=get_accelerator().current_device(),
+                      dtype=torch.float16)
+
+    kv_cache = state_manager.get_cache(0)
+
+    attn_module.build_atoms(batch)
+    if not trained_freqs:
+        out = attn_module(qkv, kv_cache, batch)
+    else:
+        inv_freqs = torch.randn((head_size // 2, ), device=get_accelerator().current_device(), dtype=torch.float16)
+        out = attn_module(qkv, kv_cache, batch, inv_freqs)
+
+    if validate_accuracy and trained_freqs is None:
+        cu_seqlens_q = torch.tensor([0] + list(itertools.accumulate([seq[0] for seq in seq_params])),
+                                    dtype=torch.int32,
+                                    device=get_accelerator().current_device())
+        cu_seqlens_kv = torch.tensor([0] + list(itertools.accumulate([seq[1] + seq[0] for seq in seq_params])),
+                                     dtype=torch.int32,
+                                     device=get_accelerator().current_device())
+
+        inflight_kv = qkv[:, head_size * n_heads_q:]
+        full_kvs = []
+        for i, kv in enumerate(kvs):
+            if kv is not None:
+                full_kvs.append(torch.cat([kv, inflight_kv[cu_seqlens_q[i]:cu_seqlens_q[i + 1]]], dim=0))
+            else:
+                full_kvs.append(inflight_kv[cu_seqlens_q[i]:cu_seqlens_q[i + 1]])
+        run_kvs = torch.cat(full_kvs, dim=0)
+        k = run_kvs[:, :head_size * n_heads_kv]
+        v = run_kvs[:, head_size * n_heads_kv:]
+
+        q = qkv[:, :head_size * n_heads_q]
+        q_ref = q.reshape((batch.current_tokens, n_heads_q, head_size))
+        k_ref = k.reshape((k.shape[0], n_heads_kv, head_size))
+        v_ref = v.reshape((v.shape[0], n_heads_kv, head_size))
+
+        max_seqlen_q = max([seq[0] for seq in seq_params])
+        max_seqlen_kv = max([seq[1] + seq[0] for seq in seq_params])
+
+        ref_o = flash_attn_varlen_func(q_ref,
+                                       k_ref,
+                                       v_ref,
+                                       cu_seqlens_q,
+                                       cu_seqlens_kv,
+                                       max_seqlen_q,
+                                       max_seqlen_kv,
+                                       softmax_scale=1.0,
+                                       causal=True)
+
+        ref_o = ref_o.reshape(batch.current_tokens, head_size * n_heads_q)
+
+        assert allclose(out, ref_o)
+
+    get_accelerator().synchronize()
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("n_tokens", [2, 33, 65, 128, 256, 2037])
+def test_single_prompt(n_tokens: int) -> None:
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(n_tokens, 0)]
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("prompt_lengths", [(128, 128), (192, 38), (514, 713), (83, 312, 610)])
+def test_multiple_prompts(prompt_lengths: Tuple[int, int]) -> None:
+    """
+    Test multiple prompts in a single batch.
+    """
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(prompt_lengths[i], 0) for i in range(len(prompt_lengths))]
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("seq_params", [(1, 34), (43, 40), (1, 144), (64, 128), (332, 628)])
+def test_continuation(seq_params: Tuple[int, int]) -> None:
+    """
+    Test continued generation/prompt processing.
+    """
+    head_size = 64
+    n_heads_q = 32
+    n_heads_kv = 32
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, [seq_params])
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("head_size", [64, 128])
+def test_head_size(head_size: int) -> None:
+    n_heads_q = 16
+    n_heads_kv = 16
+    seq_params = [(128, 128), (192, 38), (1, 814)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("head_config", [(32, 8), (64, 16), (40, 8)])
+def test_gqa(head_config: Tuple[int, int]) -> None:
+    head_size = 128
+    n_heads_q = head_config[0]
+    n_heads_kv = head_config[1]
+
+    seq_params = [(128, 128), (192, 38), (1, 814)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+def test_fully_composed() -> None:
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(332, 628), (1, 718), (1, 323), (180, 5), (224, 0)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("trained_freqs", [True, False])
+def test_rotary_emb(trained_freqs: bool) -> None:
+    head_size = 64
+    n_heads_q = 16
+    n_heads_kv = 16
+
+    seq_params = [(332, 628), (1, 718), (1, 323), (180, 5), (224, 0)]
+
+    _blocked_flash_testing_helper(head_size, n_heads_q, n_heads_kv, seq_params, trained_freqs=trained_freqs)

tests/unit/inference/v2/modules/test_cuda_pre_ln_module.py

@@ -0,0 +1,88 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional, Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSNormConfig
+from deepspeed.inference.v2.modules.interfaces import DSPreNormRegistry
+from ...v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_implementation(residual: torch.Tensor, hidden_states: Optional[torch.Tensor], gamma: torch.Tensor,
+                             beta: torch.Tensor, epsilon: float) -> Tuple[torch.Tensor, torch.Tensor]:
+    dtype = residual.dtype
+
+    residual = residual.to(torch.float32)
+    gamma = gamma.to(torch.float32)
+    beta = beta.to(torch.float32)
+
+    if hidden_states is not None:
+        hidden_states = hidden_states.to(torch.float32)
+        residual = residual + hidden_states
+    hidden_states = torch.nn.functional.layer_norm(residual, (residual.size(-1), ),
+                                                   weight=gamma,
+                                                   bias=beta,
+                                                   eps=epsilon)
+    return residual.to(dtype), hidden_states.to(dtype)
+
+
+def _pre_ln_test_helper(n_tokens: int, n_channels: int, dtype: torch.dtype, res_add: bool = False):
+    config = DSNormConfig(max_tokens=2048,
+                          type="layer_norm",
+                          channels=n_channels,
+                          residual_dtype=dtype,
+                          input_dtype=dtype,
+                          output_dtype=dtype,
+                          eps=1e-5)
+    bundle = ConfigBundle(name='cuda_pre_ln', config=config)
+
+    # Input vals
+    if res_add:
+        hidden_states = torch.randn((n_tokens, n_channels),
+                                    dtype=dtype,
+                                    device=get_accelerator().current_device_name())
+    else:
+        hidden_states = None
+
+    residual = torch.randn((n_tokens, n_channels), dtype=dtype, device=get_accelerator().current_device_name())
+    gamma = torch.randn((n_channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    beta = torch.rand((n_channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    epsilon = 1e-5
+
+    # Reference output
+    ref_residual, ref_output = reference_implementation(residual, hidden_states, gamma, beta, epsilon)
+
+    # New output
+    pre_ln_module = DSPreNormRegistry.instantiate_config(bundle)
+    gamma = pre_ln_module.transform_param(gamma)
+    beta = pre_ln_module.transform_param(beta)
+
+    ds_residual, ds_output = pre_ln_module(residual, hidden_states, gamma, beta)
+
+    # Check
+    assert allclose(ds_residual, ref_residual)
+    assert allclose(ds_output, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, channels", [(1, 2048), (37, 8192), (1280, 768), (2048, 5120)])
+def test_token_channels(tokens: int, channels: int) -> None:
+    _pre_ln_test_helper(tokens, channels, torch.float16)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("dtype", get_dtypes(include_float=False))
+def test_dtype(dtype: torch.dtype) -> None:
+    _pre_ln_test_helper(733, 2560, dtype)
+
+
+@pytest.mark.inference_v2_ops
+def test_no_res_add():
+    _pre_ln_test_helper(733, 2560, torch.float16, res_add=False)

tests/unit/inference/v2/modules/test_custom_module.py

@@ -0,0 +1,76 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.interfaces import DSPostNormRegistry
+from deepspeed.inference.v2.modules.configs import DSNormConfig
+from deepspeed.inference.v2.modules.implementations import cuda_post_ln
+from ...v2.inference_test_utils import allclose
+
+
+def reference_implementation(residual: torch.Tensor, hidden_states: torch.Tensor, gamma: torch.Tensor,
+                             beta: torch.Tensor, epsilon: float) -> torch.Tensor:
+    residual_f = residual.to(torch.float32)
+    hidden_states_f = hidden_states.to(torch.float32)
+    gamma_f = gamma.to(torch.float32)
+    beta_f = beta.to(torch.float32)
+    return torch.nn.functional.layer_norm(residual_f + hidden_states_f, (hidden_states_f.size(-1), ),
+                                          weight=gamma_f,
+                                          bias=beta_f,
+                                          eps=epsilon).to(hidden_states.dtype)
+
+
+@DSPostNormRegistry.register_module
+class CustomPostLNModule(cuda_post_ln.DSPostLNCUDAModule):
+
+    @staticmethod
+    def name():
+        return 'custom_post_ln'
+
+
+"""
+Here, we explicitly register an LN implementation outside the core deepspeed repo. This should
+validate that the registry is working as expected and we can implement modules outside the core
+repo.
+"""
+
+
+@pytest.mark.inference_v2_ops
+def test_custom_registration():
+    channels = 4096
+    dtype = torch.float16
+    tokens = 1024
+
+    config = DSNormConfig(max_tokens=2048,
+                          type="layer_norm",
+                          channels=channels,
+                          residual_dtype=dtype,
+                          input_dtype=dtype,
+                          output_dtype=dtype,
+                          eps=1e-5)
+    bundle = ConfigBundle(name='custom_post_ln', config=config)
+
+    # Input vals
+    hidden_states = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    residual = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    gamma = torch.randn((channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    beta = torch.rand((channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    epsilon = 1e-5
+
+    # Reference output
+    ref_output = reference_implementation(residual, hidden_states, gamma, beta, epsilon)
+
+    # New output
+    post_ln_module = DSPostNormRegistry.instantiate_config(bundle)
+    gamma = post_ln_module.transform_param(gamma)
+    beta = post_ln_module.transform_param(beta)
+    ds_output, _ = post_ln_module(residual, hidden_states, gamma, beta)
+
+    # Check
+    assert allclose(ds_output, ref_output)

tests/unit/inference/v2/modules/test_cutlass_moe.py

@@ -0,0 +1,328 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSMoEConfig
+from deepspeed.inference.v2.modules.interfaces import DSMoERegistry
+
+from ..kernels.ragged_ops.ragged_testing_utils import build_simple_batch
+from ...v2.inference_test_utils import allclose, get_dtypes
+
+
+def _gating_reference(logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Reference gating code.
+    """
+    logits = logits.float()
+    probs = torch.nn.functional.softmax(logits, dim=1)
+
+    indices1_s = torch.argmax(probs, dim=-1)
+    mask1 = torch.nn.functional.one_hot(indices1_s, num_classes=logits.shape[-1])
+    indices_mask = mask1.sum(dim=1) * logits.shape[-1] - 1
+    indices1_s = torch.min(indices1_s, indices_mask)
+
+    gates1_s = (probs * mask1).sum(dim=1)
+
+    sorted_indices = indices1_s.sort()[1]
+    original_indices = sorted_indices.sort()[1]
+
+    exp_count = torch.bincount(indices1_s, minlength=logits.shape[-1]).long()
+    exp_count_cumsum = exp_count.cumsum(dim=0)
+
+    return sorted_indices, original_indices, exp_count_cumsum, gates1_s
+
+
+def _reference_impl(hidden_states: torch.Tensor, gate_weight: torch.Tensor, mlp_1_w: torch.Tensor,
+                    mlp_2_w: torch.Tensor, mlp_1_b: torch.Tensor, mlp_2_b: torch.Tensor,
+                    act_fn: ActivationType) -> torch.Tensor:
+    """
+    Reference implementation of the MoE module.
+    """
+
+    act_fn_dict = {
+        ActivationType.GELU: torch.nn.functional.gelu,
+        ActivationType.RELU: torch.nn.functional.relu,
+        ActivationType.SILU: torch.nn.functional.silu,
+        ActivationType.IDENTITY: lambda x: x,
+    }
+
+    logits = torch.matmul(hidden_states, gate_weight.t())
+    sorted_indices, original_indices, exp_count_cumsum, gate_scales = _gating_reference(logits)
+
+    moe_input = hidden_states[sorted_indices]
+
+    output_unordered = torch.empty_like(hidden_states)
+
+    for expert_idx in range(mlp_1_w.shape[0]):
+        min_bound = 0 if expert_idx == 0 else exp_count_cumsum[expert_idx - 1]
+        max_bound = exp_count_cumsum[expert_idx]
+
+        input_slice = moe_input[min_bound:max_bound]
+        intermediate = torch.nn.functional.linear(input_slice, mlp_1_w[expert_idx], mlp_1_b[expert_idx])
+
+        intermediate = act_fn_dict[act_fn](intermediate)
+        output_slice = torch.nn.functional.linear(intermediate, mlp_2_w[expert_idx], mlp_2_b[expert_idx])
+
+        output_unordered[min_bound:max_bound] = output_slice
+
+    output = output_unordered[original_indices]
+
+    output.mul_(gate_scales.unsqueeze(-1)).reshape(hidden_states.shape)
+    return output
+
+
+def _cutlass_moe_testing_helper(tokens: int,
+                                in_channels: int,
+                                intermediate_dim: int,
+                                experts: int,
+                                dtype: int,
+                                activation_type: ActivationType = ActivationType.GELU,
+                                use_bias: bool = True,
+                                iters: int = 1) -> None:
+
+    config = DSMoEConfig(max_tokens=4096,
+                         model_dim=in_channels,
+                         intermediate_features=intermediate_dim,
+                         n_experts=experts,
+                         activation=activation_type,
+                         input_dtype=dtype,
+                         output_dtype=dtype)
+
+    implementation_config = {"weight_dtype": DtypeEnum(dtype)}
+
+    bundle = ConfigBundle(name='cutlass_multi_gemm_moe', config=config, implementation_config=implementation_config)
+    moe_module = DSMoERegistry.instantiate_config(bundle)
+
+    batch = build_simple_batch([tokens])
+
+    # Parameters
+    gate_weight = torch.randn(
+        (experts, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+
+    mlp_1_w = torch.randn(
+        (experts, intermediate_dim, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+    mlp_2_w = torch.randn(
+        (experts, in_channels, intermediate_dim), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+
+    if use_bias:
+        mlp_1_b = torch.randn(
+            (experts, intermediate_dim), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+        mlp_2_b = torch.randn(
+            (experts, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+    else:
+        mlp_1_b = None
+        mlp_2_b = None
+
+    gate_ds = moe_module.transform_gate_param(gate_weight)
+    mlp_1_w_ds = moe_module.transform_moe_mlp_1_param(mlp_1_w)
+    mlp_1_b_ds = moe_module.transform_moe_mlp_1_param(mlp_1_b)
+    mlp_2_w_ds = moe_module.transform_moe_mlp_2_param(mlp_2_w)
+    mlp_2_b_ds = moe_module.transform_moe_mlp_2_param(mlp_2_b)
+
+    for _ in range(iters):
+        # Input vals
+        hidden_states = torch.randn(
+            (tokens, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+
+        # Reference implementation
+        ref_output = _reference_impl(hidden_states, gate_weight, mlp_1_w, mlp_2_w, mlp_1_b, mlp_2_b, activation_type)
+
+        output = moe_module(hidden_states,
+                            batch,
+                            gate_ds,
+                            mlp_1_w_ds,
+                            mlp_2_w_ds,
+                            mlp_1_b=mlp_1_b_ds,
+                            mlp_2_b=mlp_2_b_ds)
+
+        # Increase the tolerance for larger meta ops since the error is additive
+        assert allclose(output, ref_output, tolerances=(1e-2, 1e-2))
+
+    get_accelerator().synchronize()
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("experts", [2, 32, 64])
+def test_expert_variance(experts: int) -> None:
+    _cutlass_moe_testing_helper(tokens=876,
+                                in_channels=4096,
+                                intermediate_dim=2048,
+                                experts=experts,
+                                dtype=DtypeEnum.fp16,
+                                activation_type=ActivationType.IDENTITY,
+                                use_bias=True)
+
+
+@pytest.mark.inference_v2_ops
+def test_successive_inputs():
+    """
+    The CUTLASS MoE uses persistent state (expert counts) that is assumed to be cleared
+    on each forward pass. This ensures that the module is clearing that metadata.
+    """
+    _cutlass_moe_testing_helper(tokens=876,
+                                in_channels=4096,
+                                intermediate_dim=2048,
+                                experts=64,
+                                dtype=DtypeEnum.fp16,
+                                activation_type=ActivationType.IDENTITY,
+                                use_bias=True,
+                                iters=10)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("dtype", get_dtypes(include_float=False))
+def test_dtypes(dtype: torch.dtype) -> None:
+    _cutlass_moe_testing_helper(tokens=876,
+                                in_channels=4096,
+                                intermediate_dim=2048,
+                                experts=64,
+                                dtype=DtypeEnum(dtype),
+                                activation_type=ActivationType.IDENTITY,
+                                use_bias=True)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("activation_type", [ActivationType.GELU, ActivationType.RELU, ActivationType.SILU])
+def test_activation_types(activation_type: ActivationType) -> None:
+    _cutlass_moe_testing_helper(tokens=876,
+                                in_channels=4096,
+                                intermediate_dim=2048,
+                                experts=64,
+                                dtype=DtypeEnum.fp16,
+                                activation_type=activation_type,
+                                use_bias=True)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("in_channels, out_channels", [(4096, 2048), (2048, 8192), (6144, 3072)])
+def test_in_out_channels(in_channels: int, out_channels: int) -> None:
+    _cutlass_moe_testing_helper(tokens=876,
+                                in_channels=in_channels,
+                                intermediate_dim=out_channels,
+                                experts=64,
+                                dtype=DtypeEnum.fp16,
+                                activation_type=ActivationType.IDENTITY,
+                                use_bias=True)
+
+
+def _mixtral_moe_baseline(hidden_states: torch.Tensor,
+                          gate_weight: torch.Tensor,
+                          mlp_w1: torch.Tensor,
+                          mlp_w2: torch.Tensor,
+                          mlp_w3: torch.Tensor,
+                          force_float: bool = False) -> torch.Tensor:
+    """
+    Baseline implementation for mixtral MoE module.
+
+    Based on transformers implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/models/mixtral/modeling_mixtral.py
+    """
+    output_dtype = hidden_states.dtype
+    if force_float:
+        hidden_states = hidden_states.float()
+        gate_weight = gate_weight.float()
+        mlp_w1 = mlp_w1.float()
+        mlp_w2 = mlp_w2.float()
+        mlp_w3 = mlp_w3.float()
+
+    router_logits = torch.nn.functional.linear(hidden_states, gate_weight)
+    routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)
+    routing_weights, selected_experts = routing_weights.topk(k=2, dim=-1)
+    routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+
+    # NOTE(cmikeh2): This is a difference implementation, ours will preserve the original scale
+    # as float32 and perform in-kernel fused FP16->FP32->FP16 conversion.
+    routing_weights = routing_weights.to(hidden_states.dtype)
+
+    final_hidden_states = torch.zeros_like(hidden_states)
+
+    expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=gate_weight.shape[0]).permute(2, 1, 0)
+    get_accelerator().synchronize()
+
+    for expert_idx in range(gate_weight.shape[0]):
+        exp_mlp_w1 = mlp_w1[expert_idx]
+        exp_mlp_w2 = mlp_w2[expert_idx]
+        exp_mlp_w3 = mlp_w3[expert_idx]
+
+        idx, top_x = torch.where(expert_mask[expert_idx])
+
+        if top_x.shape[0] == 0:
+            continue
+
+        top_x_list = top_x.tolist()
+        idx_list = idx.tolist()
+
+        current_state = hidden_states[top_x_list]
+
+        linear = torch.nn.functional.linear
+        intermediate = torch.nn.functional.silu(linear(current_state, exp_mlp_w1)) * linear(current_state, exp_mlp_w3)
+        output = linear(intermediate, exp_mlp_w2) * routing_weights[top_x_list, idx_list].unsqueeze(-1)
+        final_hidden_states.index_add_(0, top_x, output.to(final_hidden_states.dtype))
+
+    return final_hidden_states.to(output_dtype)
+
+
+@pytest.mark.inference_v2_ops
+def test_mixtral_moe_config():
+
+    experts = 8
+    n_top_k = 2
+    in_channels = 4096
+    intermediate_dim = 2048
+    dtype = DtypeEnum.bf16
+
+    # Parameters
+    gate_weight = torch.randn(
+        (experts, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+
+    mlp_w1 = torch.randn(
+        (experts, intermediate_dim, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+    mlp_w3 = torch.randn(
+        (experts, intermediate_dim, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+    mlp_w2 = torch.randn(
+        (experts, in_channels, intermediate_dim), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+
+    n_tokens = 256
+    hidden_states = torch.randn(
+        (n_tokens, in_channels), dtype=dtype.value, device=get_accelerator().current_device()) * .1
+
+    baseline = _mixtral_moe_baseline(hidden_states, gate_weight, mlp_w1, mlp_w2, mlp_w3)
+
+    mlp_w13_fused = torch.cat([mlp_w1, mlp_w3], dim=-1).reshape(experts, 2 * intermediate_dim, in_channels)
+
+    config = DSMoEConfig(max_tokens=4096,
+                         model_dim=in_channels,
+                         intermediate_features=intermediate_dim,
+                         n_experts=experts,
+                         activation=ActivationType.SiGLU,
+                         input_dtype=dtype,
+                         output_dtype=dtype,
+                         top_k=n_top_k,
+                         normalize_scores=True)
+
+    implementation_config = {"weight_dtype": DtypeEnum(dtype)}
+
+    bundle = ConfigBundle(name='cutlass_multi_gemm_moe', config=config, implementation_config=implementation_config)
+    moe_module = DSMoERegistry.instantiate_config(bundle)
+
+    batch = build_simple_batch([n_tokens])
+
+    gate_ds = moe_module.transform_gate_param(gate_weight)
+    mlp_w1_ds = moe_module.transform_moe_mlp_1_param(mlp_w13_fused)
+    mlp_w2_ds = moe_module.transform_moe_mlp_2_param(mlp_w2)
+
+    output = moe_module(hidden_states, batch, gate_ds, mlp_w1_ds, mlp_w2_ds)
+
+    # NOTE(cmikeh2): These are higher than the other tests for reasons that aren't quite
+    # clear to me. My best guess is that the SiGLU activation is causing larger numerical
+    # divergence. The thresholds chosen here is based on the observed error between the
+    # float and bfloat16 reference implementations.
+    assert allclose(output, baseline.to(dtype.value), tolerances=(5e-2, 5e-2))

tests/unit/inference/v2/modules/test_post_ln_module.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSNormConfig
+from deepspeed.inference.v2.modules.interfaces import DSPostNormRegistry
+from ...v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_implementation(residual: torch.Tensor, hidden_states: torch.Tensor, gamma: torch.Tensor,
+                             beta: torch.Tensor, epsilon: float) -> torch.Tensor:
+    residual_f = residual.to(torch.float32)
+    hidden_states_f = hidden_states.to(torch.float32)
+    gamma_f = gamma.to(torch.float32)
+    beta_f = beta.to(torch.float32)
+    return torch.nn.functional.layer_norm(residual_f + hidden_states_f, (hidden_states_f.size(-1), ),
+                                          weight=gamma_f,
+                                          bias=beta_f,
+                                          eps=epsilon).to(hidden_states.dtype)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, channels", [(1, 2048), (37, 8192), (1280, 768), (2048, 5120)])
+@pytest.mark.parametrize("dtype", get_dtypes())
+def test_cuda_post_ln_module(tokens: int, channels: int, dtype: torch.dtype) -> None:
+    config = DSNormConfig(max_tokens=2048,
+                          type="layer_norm",
+                          channels=channels,
+                          residual_dtype=dtype,
+                          input_dtype=dtype,
+                          output_dtype=dtype,
+                          eps=1e-5)
+    bundle = ConfigBundle(name='cuda_post_ln', config=config)
+
+    # Input vals
+    hidden_states = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    residual = torch.randn((tokens, channels), dtype=dtype, device=get_accelerator().current_device_name())
+    gamma = torch.randn((channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    beta = torch.rand((channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    epsilon = 1e-5
+
+    # Reference output
+    ref_output = reference_implementation(residual, hidden_states, gamma, beta, epsilon)
+
+    # New output
+    post_ln_module = DSPostNormRegistry.instantiate_config(bundle)
+    gamma = post_ln_module.transform_param(gamma)
+    beta = post_ln_module.transform_param(beta)
+    ds_output, _ = post_ln_module(residual, hidden_states, gamma, beta)
+
+    # Check
+    assert allclose(ds_output, ref_output)

tests/unit/inference/v2/modules/test_pre_rms_module.py

@@ -0,0 +1,88 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional, Tuple
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSNormConfig
+from deepspeed.inference.v2.modules.interfaces import DSPreNormRegistry
+from ...v2.inference_test_utils import get_dtypes, allclose
+
+
+def reference_implementation(residual: torch.Tensor, hidden_states: Optional[torch.Tensor], gamma: torch.Tensor,
+                             epsilon: float) -> Tuple[torch.Tensor, torch.Tensor]:
+    dtype = residual.dtype
+
+    if hidden_states is not None:
+        hidden_states = hidden_states
+        residual = residual + hidden_states
+
+    rms_vals = residual.to(torch.float32)
+    variance = rms_vals.pow(2).mean(-1, keepdim=True)
+    rms_vals = rms_vals * torch.rsqrt(variance + epsilon)
+
+    if gamma.dtype in [torch.float16, torch.bfloat16]:
+        rms_vals = rms_vals.to(gamma.dtype)
+
+    hidden_states = gamma * rms_vals
+
+    return residual.to(dtype), hidden_states.to(dtype)
+
+
+def _pre_rms_test_helper(n_tokens: int, n_channels: int, dtype: torch.dtype, res_add: bool = False):
+    config = DSNormConfig(max_tokens=2048,
+                          type="rms_norm",
+                          channels=n_channels,
+                          residual_dtype=dtype,
+                          input_dtype=dtype,
+                          output_dtype=dtype,
+                          eps=1e-5)
+    bundle = ConfigBundle(name='cuda_pre_rms', config=config)
+
+    # Input vals
+    if res_add:
+        hidden_states = torch.randn((n_tokens, n_channels),
+                                    dtype=dtype,
+                                    device=get_accelerator().current_device_name())
+    else:
+        hidden_states = None
+
+    residual = torch.randn((n_tokens, n_channels), dtype=dtype, device=get_accelerator().current_device_name())
+    gamma = torch.randn((n_channels), dtype=torch.float32, device=get_accelerator().current_device_name())
+    epsilon = 1e-5
+
+    # Reference output
+    ref_residual, ref_output = reference_implementation(residual, hidden_states, gamma, epsilon)
+
+    # New output
+    pre_ln_module = DSPreNormRegistry.instantiate_config(bundle)
+    gamma = pre_ln_module.transform_param(gamma)
+
+    ds_residual, ds_output = pre_ln_module(residual, hidden_states, gamma)
+
+    # Check
+    assert allclose(ds_residual, ref_residual)
+    assert allclose(ds_output, ref_output)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens, channels", [(1, 2048), (37, 8192), (1280, 768), (2048, 5120)])
+def test_token_channels(tokens: int, channels: int) -> None:
+    _pre_rms_test_helper(tokens, channels, torch.float16)
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("dtype", get_dtypes(include_float=False))
+def test_dtype(dtype: torch.dtype) -> None:
+    _pre_rms_test_helper(733, 2560, dtype)
+
+
+@pytest.mark.inference_v2_ops
+def test_no_res_add():
+    _pre_rms_test_helper(733, 2560, torch.float16, res_add=False)

tests/unit/inference/v2/modules/test_quantized_linear_module.py

@@ -0,0 +1,183 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import Optional
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.inference_utils import ActivationType, DtypeEnum, is_gated
+from deepspeed.inference.v2.modules import ConfigBundle
+from deepspeed.inference.v2.modules.configs import DSLinearConfig
+from deepspeed.inference.v2.modules.interfaces import DSLinearRegistry
+from ...v2.inference_test_utils import allclose
+
+
+def reference_implementation(hidden_states: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor],
+                             act_type: ActivationType) -> torch.Tensor:
+    dtype = hidden_states.dtype
+    out_states = torch.nn.functional.linear(hidden_states, weight, bias)
+    out_states.float()
+
+    if is_gated(act_type):
+        act_func_map = {
+            ActivationType.ReGLU: torch.nn.functional.relu,
+            ActivationType.GEGLU: lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
+            ActivationType.SiGLU: torch.nn.functional.silu,
+        }
+
+        act_act = out_states[..., ::2]
+        act_linear = out_states[..., 1::2]
+
+        act_act = act_func_map[act_type](act_act)
+        out_states = act_act * act_linear
+    else:
+        act_func_map = {
+            ActivationType.RELU: torch.nn.functional.relu,
+            ActivationType.GELU: torch.nn.functional.gelu,
+            ActivationType.SILU: torch.nn.functional.silu,
+            ActivationType.IDENTITY: lambda x: x,
+        }
+
+        out_states = act_func_map[act_type](out_states)
+    return out_states.to(dtype)
+
+
+def _fp6_quant_dequant_weights(weight: torch.Tensor) -> torch.Tensor:
+    from deepspeed.inference.v2.modules.implementations.linear.quantized_linear import fp_quantize
+    weight_quantized_fake_fp6, scales = fp_quantize(weight, num_bits=6, exp_bits=3)
+    return weight_quantized_fake_fp6 * scales
+
+
+def quant_dequant_implementation(hidden_states: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor],
+                                 act_type: ActivationType) -> torch.Tensor:
+    dtype = hidden_states.dtype
+    weight_dequantized = _fp6_quant_dequant_weights(weight)
+    out_states = torch.nn.functional.linear(hidden_states, weight_dequantized, bias)
+    out_states.float()
+
+    if is_gated(act_type):
+        act_func_map = {
+            ActivationType.ReGLU: torch.nn.functional.relu,
+            ActivationType.GEGLU: lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
+            ActivationType.SiGLU: torch.nn.functional.silu,
+        }
+
+        act_act = out_states[..., ::2]
+        act_linear = out_states[..., 1::2]
+
+        act_act = act_func_map[act_type](act_act)
+        out_states = act_act * act_linear
+    else:
+        act_func_map = {
+            ActivationType.RELU: torch.nn.functional.relu,
+            ActivationType.GELU: torch.nn.functional.gelu,
+            ActivationType.SILU: torch.nn.functional.silu,
+            ActivationType.IDENTITY: lambda x: x,
+        }
+
+        out_states = act_func_map[act_type](out_states)
+    return out_states.to(dtype)
+
+
+def _fp6_quantized_linear_helper(tokens: int,
+                                 in_channels: int,
+                                 out_channels: int,
+                                 dtype: DtypeEnum,
+                                 act_fn: ActivationType,
+                                 use_bias: bool = True,
+                                 expect_failure: bool = False) -> None:
+    # The current FP6 kernel only supports NVIDIA Ampere GPUs.
+    if not 'cuda' in get_accelerator().current_device_name():
+        return
+    major, _ = torch.cuda.get_device_capability()  #ignore-cuda
+    if major != 8:
+        return
+
+    # Input vals
+    hidden_states = torch.randn(
+        (tokens, in_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
+
+    weight_out_channels = 2 * \
+        out_channels if is_gated(act_fn) else out_channels
+    weight = torch.randn(
+        (weight_out_channels, in_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
+    if use_bias:
+        bias = torch.randn(
+            (weight_out_channels), dtype=dtype.value, device=get_accelerator().current_device_name()) * .01
+    else:
+        bias = None
+
+    # quantize and dequantize output
+    ref_quant_dequant_output = quant_dequant_implementation(hidden_states, weight, bias, act_fn)
+
+    linear_config = DSLinearConfig(max_tokens=2048,
+                                   in_channels=in_channels,
+                                   out_channels=out_channels,
+                                   activation=act_fn,
+                                   input_dtype=dtype,
+                                   output_dtype=dtype)
+    bundle = ConfigBundle(name='quantized_wf6af16_linear', config=linear_config)
+    fp6_linear_module = DSLinearRegistry.instantiate_config(bundle)
+    weight_fp6 = fp6_linear_module.transform_param(weight.clone().cpu()).to(get_accelerator().current_device_name())
+
+    if expect_failure:
+        with pytest.raises(ValueError) as excinfo:
+            ds_output = fp6_linear_module(hidden_states, weight_fp6, bias)
+        assert "The out and in channel should be multiple of 256 and 64 respectively." in str(excinfo.value)
+    else:
+        ds_output = fp6_linear_module(hidden_states, weight_fp6, bias)
+        # The current FP6 kernel uses FP16 Tensor Core.
+        tolerances = (3e-2, 2e-3)  # tolerances for fp16
+
+        # Check DeepSpeed implementation
+        assert allclose(ds_output, ref_quant_dequant_output, tolerances=tolerances)
+
+
+all_acts = [
+    ActivationType.RELU,
+    ActivationType.GELU,
+    ActivationType.SILU,
+    ActivationType.GEGLU,
+    ActivationType.ReGLU,
+    ActivationType.SiGLU,
+]
+all_tokens = [37]
+all_in_out_channels = [
+    (4096, 4096),
+]
+
+
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens", all_tokens)
+@pytest.mark.parametrize("in_channels, out_channels", all_in_out_channels)
+@pytest.mark.parametrize("act_fn", all_acts)
+@pytest.mark.parametrize("use_bias", [True, False])
+def test_fp6_quantized_linear_act_fn(tokens: int, in_channels: int, out_channels: int, act_fn: ActivationType,
+                                     use_bias: bool) -> None:
+    _fp6_quantized_linear_helper(tokens=tokens,
+                                 in_channels=in_channels,
+                                 out_channels=out_channels,
+                                 dtype=DtypeEnum.fp16,
+                                 act_fn=act_fn,
+                                 use_bias=use_bias)
+
+
+# Other shapes, not supported by FP6 kernels. Will raise ValueError.
+@pytest.mark.inference_v2_ops
+@pytest.mark.parametrize("tokens", all_tokens)
+@pytest.mark.parametrize("in_channels, out_channels", [(4608, 1728)])
+@pytest.mark.parametrize("act_fn", all_acts)
+@pytest.mark.parametrize("use_bias", [True, False])
+def test_fp6_quantized_linear_act_fn_fail(tokens: int, in_channels: int, out_channels: int, act_fn: ActivationType,
+                                          use_bias: bool) -> None:
+    _fp6_quantized_linear_helper(tokens=tokens,
+                                 in_channels=in_channels,
+                                 out_channels=out_channels,
+                                 dtype=DtypeEnum.fp16,
+                                 act_fn=act_fn,
+                                 use_bias=use_bias,
+                                 expect_failure=True)

tests/unit/inference/v2/ragged/__init__.py

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

tests/unit/inference/v2/ragged/test_blocked_allocator.py

@@ -0,0 +1,166 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import random
+from typing import List
+
+import pytest
+import torch
+
+from deepspeed.inference.v2.ragged.blocked_allocator import BlockedAllocator
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('bad_size', [0, -1])
+def test_bad_initialization(bad_size: int) -> None:
+    with pytest.raises(ValueError):
+        BlockedAllocator(bad_size)
+
+
+@pytest.mark.inference_v2
+def test_allocation() -> None:
+
+    allocator = BlockedAllocator(16)
+
+    a1 = allocator.allocate(4)
+    assert a1.numel() == 4
+    assert allocator.free_blocks == 12
+
+    a2_allocs = []
+    for i in range(3):
+        a2_allocs.append(allocator.allocate(2))
+        assert allocator.free_blocks == 12 - (i + 1) * 2
+
+    a3 = allocator.allocate(6)
+    assert a3.numel() == 6
+
+    assert allocator.free_blocks == 0
+
+    # Test that we can't allocate more blocks than we have.
+    with pytest.raises(ValueError):
+        allocator.allocate(1)
+
+    all_vals = torch.cat([a1, *a2_allocs, a3], dim=0)
+    unique_vals = torch.unique(all_vals, sorted=False)
+    assert unique_vals.numel() == all_vals.numel()
+
+
+@pytest.mark.inference_v2
+def test_too_large_allocation():
+    allocator = BlockedAllocator(16)
+
+    with pytest.raises(ValueError):
+        allocator.allocate(17)
+
+
+@pytest.mark.inference_v2
+def test_deallocation() -> None:
+    allocator = BlockedAllocator(16)
+
+    # Allocate
+    all_blocks = allocator.allocate(16)
+    assert allocator.free_blocks == 0
+
+    # Deallocate all blocks
+    allocator.free(all_blocks)
+    assert allocator.free_blocks == 16
+
+    # Get all the blocks again
+    all_blocks = allocator.allocate(16)
+
+    # Deallocate in chunks
+    c1 = all_blocks[:4]
+    c2 = all_blocks[4:8]
+
+    allocator.free(c1)
+    assert allocator.free_blocks == 4
+
+    allocator.free(c2)
+    assert allocator.free_blocks == 8
+
+    with pytest.raises(ValueError):
+        allocator.free(c1)
+
+    with pytest.raises(ValueError):
+        allocator.free(c2)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('index', [-1, 2])
+def test_invalid_dealloc_indices(index: int):
+    allocator = BlockedAllocator(1)
+
+    with pytest.raises(ValueError):
+        allocator.free(torch.tensor([index]))
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('index', [-1, 2])
+def test_invalid_alloc_indices(index: int):
+    allocator = BlockedAllocator(1)
+    allocator.allocate(1)
+
+    to_free = [0, index]
+
+    with pytest.raises(ValueError):
+        allocator.free(torch.tensor(to_free))
+
+    # Block 0 should not be freed if passed with an invalid index.
+    assert allocator.free_blocks == 0
+
+    allocator.free(torch.tensor([0]))
+    assert allocator.free_blocks == 1
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('test_iters', [8192])
+def test_long_running_allocation(test_iters: int) -> None:
+    """
+    Evaluate the stability of the allocator over a longer sequence of allocations/deallocations.
+    """
+    TOTAL_BLOCKS = 128
+
+    allocator = BlockedAllocator(TOTAL_BLOCKS)
+
+    def validate_uniqueness(all_blocks: List[torch.Tensor]) -> None:
+        all_vals = torch.cat(all_blocks, dim=0)
+        assert all_vals.numel() <= TOTAL_BLOCKS
+
+        unique_vals = torch.unique(all_vals, sorted=False)
+        assert unique_vals.numel() == all_vals.numel()
+
+    all_allocs: List[torch.Tensor] = []
+    num_allocs = 0
+    num_frees = 0
+    num_blocks_allocated = 0
+    num_blocks_freed = 0
+
+    for _ in range(test_iters):
+        decision = random.randint(0, 1)
+
+        if decision == 0:
+            blocks_to_allocate = random.randint(1, 24)
+            if blocks_to_allocate > allocator.free_blocks:
+                with pytest.raises(ValueError):
+                    allocator.allocate(blocks_to_allocate)
+            else:
+                all_allocs.append(allocator.allocate(blocks_to_allocate))
+                num_allocs += 1
+                num_blocks_allocated += blocks_to_allocate
+        else:
+            if len(all_allocs) > 0:
+                idx = random.randint(0, len(all_allocs) - 1)
+                allocator.free(all_allocs[idx])
+
+                num_frees += 1
+                num_blocks_freed += all_allocs[idx].numel()
+
+                del all_allocs[idx]
+
+        if len(all_allocs) > 0:
+            validate_uniqueness(all_allocs)
+
+    assert num_allocs == num_frees + len(all_allocs)
+    assert num_blocks_allocated == num_blocks_freed + (TOTAL_BLOCKS - allocator.free_blocks)

tests/unit/inference/v2/ragged/test_manager_configs.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import pytest
+
+from pydantic import ValidationError
+
+from deepspeed.inference.v2.ragged import DSStateManagerConfig
+
+
+@pytest.mark.inference_v2
+def test_negative_max_tracked_sequences() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_tracked_sequences=-1)
+
+
+@pytest.mark.inference_v2
+def test_zero_max_tracked_sequences() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_tracked_sequences=0)
+
+
+@pytest.mark.inference_v2
+def test_negative_max_ragged_batch_size() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_ragged_batch_size=-1)
+
+
+@pytest.mark.inference_v2
+def test_zero_max_ragged_batch_size() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_ragged_batch_size=0)
+
+
+@pytest.mark.inference_v2
+def test_negative_max_ragged_sequence_count() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_ragged_sequence_count=-1)
+
+
+@pytest.mark.inference_v2
+def test_zero_max_ragged_sequence_count() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_ragged_sequence_count=0)
+
+
+@pytest.mark.inference_v2
+def test_too_small_max_ragged_batch_size() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_ragged_batch_size=512, max_ragged_sequence_count=1024)
+
+
+@pytest.mark.inference_v2
+def test_too_small_max_tracked_sequences() -> None:
+    with pytest.raises(ValidationError):
+        DSStateManagerConfig(max_tracked_sequences=512, max_ragged_sequence_count=1024)

tests/unit/inference/v2/ragged/test_ragged_wrapper.py

@@ -0,0 +1,112 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from typing import List
+
+import pytest
+import torch
+
+from deepspeed.accelerator import get_accelerator
+from deepspeed.inference.v2.ragged import (
+    PlaceholderSequenceDescriptor,
+    RaggedBatchWrapper,
+    DSStateManagerConfig,
+)
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('max_ragged_sequence_count, max_ragged_batch_size', [(128, 512), (128, 1024)])
+def test_wrapper_initialization(max_ragged_sequence_count: int, max_ragged_batch_size: int) -> None:
+    config = DSStateManagerConfig(max_tracked_sequences=max_ragged_sequence_count,
+                                  max_ragged_batch_size=max_ragged_batch_size,
+                                  max_ragged_sequence_count=max_ragged_sequence_count)
+
+    batch = RaggedBatchWrapper(config)
+
+    assert batch.current_tokens == 0
+    assert batch.current_sequences == 0
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('seq_len', [1, 37, 128, 512])
+def test_single_sequence_batch(seq_len: int) -> None:
+    """
+    Test we successfully construct single sequence batches and the on device metadata is accurate.
+    """
+
+    config = DSStateManagerConfig()
+    batch = RaggedBatchWrapper(config)
+
+    batch.clear()
+
+    assert batch.current_tokens == 0
+    assert batch.current_sequences == 0
+
+    seq_desc = PlaceholderSequenceDescriptor()
+    tokens = torch.randint(0, 100, (seq_len, ))
+    batch.insert_sequence(seq_desc, tokens)
+
+    batch.finalize()
+
+    assert batch.current_tokens == seq_len
+    assert batch.current_sequences == 1
+    assert torch.equal(batch.input_ids(), tokens.to(get_accelerator().current_device()))
+    assert torch.equal(batch.tokens_to_seq(), torch.zeros_like(tokens, device=get_accelerator().current_device()))
+    assert torch.equal(batch.batch_metadata_buffer(),
+                       torch.tensor([seq_len, 1], device=get_accelerator().current_device()))
+
+    batch.clear()
+
+    assert batch.current_tokens == 0
+    assert batch.current_sequences == 0
+
+
+@pytest.mark.inference_v2
+@pytest.mark.parametrize('seq_lens', [[128, 128], [1, 32, 243], [64, 1, 1, 1, 1, 393, 27, 2]])
+def test_multi_sequence_batch(seq_lens: List[int]) -> None:
+    """
+    Test sequentially adding new tokens to a batch and validate device data structures hold
+    the appropriate data.
+    """
+    config = DSStateManagerConfig()
+    batch = RaggedBatchWrapper(config)
+
+    batch.clear()
+
+    assert batch.current_tokens == 0
+    assert batch.current_sequences == 0
+
+    all_toks = [torch.randint(0, 100, (seq_len, )) for seq_len in seq_lens]
+
+    for i, toks in enumerate(all_toks):
+        seq_desc = PlaceholderSequenceDescriptor()
+        batch.insert_sequence(seq_desc, toks)
+
+        assert batch.current_tokens == sum(seq_lens[:i + 1])
+        assert batch.current_sequences == i + 1
+
+    batch.finalize()
+
+    assert batch.current_tokens == sum(seq_lens)
+    assert batch.current_sequences == len(seq_lens)
+
+    assert torch.equal(batch.input_ids(), torch.cat(all_toks, dim=0).to(get_accelerator().current_device()))
+    assert torch.equal(
+        batch.tokens_to_seq(),
+        torch.cat([torch.full((seq_len, ), i, dtype=torch.int32) for i, seq_len in enumerate(seq_lens)],
+                  dim=0).to(get_accelerator().current_device()))
+
+    for i, seq_len in enumerate(seq_lens):
+        assert batch.inflight_seq_descriptors()[i][0] == sum(seq_lens[:i])
+        assert batch.inflight_seq_descriptors()[i][1] == seq_len
+        assert batch.inflight_seq_descriptors()[i][2] == 0
+
+    assert torch.equal(batch.batch_metadata_buffer(),
+                       torch.tensor([sum(seq_lens), len(seq_lens)], device=get_accelerator().current_device()))
+
+    batch.clear()
+
+    assert batch.current_tokens == 0
+    assert batch.current_sequences == 0