User-facing configuration parsed from DS config JSON.

Sets parameters for ZeRO optimizations.

A config bundle is a collection of configs that are used to instantiate a model implementation.

Initialize the DeepSpeed Transformer Config. Arguments: hidden_size: The hidden size of the transformer layer intermed

CUDA implementation of gated activation kernel. This kernel assumes that the input tensor has gate and activation values in adjacent channels

A throwaway git repo seeded with BASELINE, committed on ``master``.

Configuration for Automatic Tensor Parallelism. Example usage: config = AutoTPConfig( tp_size=4, layer_specs

Unified specification for tensor parallel layer partitioning. This is inspired by Universal Checkpointing's SubparamShape but extended f

Automatic Expert Parallelism: detect and replace MoE layers.

For float16 configuration

Allocator class for managing which blocks are free/used in the blocked KV-cache. This is a simple allocator that uses a linked list to ke

Config class for DSLinearBase.

Lightweight metadata describing what an optimization pass expects and produces. Contracts let DeepCompile validate a pass schedule before it runs

Configuration settings for quantization for LoRAOptimizedLinear, QuantizedLinear, and QuantizedParameter Attributes: q_bits (int

Column-parallel linear layer with sub-parameter support. Handles cases where weights contain multiple logical sub-parameters that need t

Config class for both DSPreLN and DSPostLN.

Modules to be parallelized with pipeline parallelism. The key constraint that enables pipeline parallelism is the representation of the forwa

Initialize the DeepSpeed GPT Transformer Layer.

The DummySequenceDescriptor is an empty object that allows us to perform schedulability checks before formally tracking a sequence.

Quantized parameter class that implements weight quantization. Weights are stored in quantized form on GPUs, and can be dequantized on-the-fl

Container for all the auxiliary Tensors used in the management of a ragged batch. For each Tensor, we maintain a shadow Tensor on the host.

CUDA implementation of top-1 gating. This will perform a softmax on the logits, and return the scale as well as its idx within that expert's

Group of timers. Borrowed from Nvidia Megatron code

CUDA implementation of bias activation kernel. This kernel should be deprecated once we are fusing the bias activation into the linear kernel

Initialize the DeepSpeed Transformer Config. Arguments: batch_size: The maximum batch size used for running the kernel on each GP

InternVL-style splice: replace IMG_CONTEXT token runs with visual tokens. InternVL encodes each image as ``<IMG_START> <IMG_CONTEXT>×N <IMG_END>`

CUDA Kernel implementation that will perform rotary position embeddings on the queries and keys before copying into a blocked KV cache.

Initialize an MoE layer. Arguments: hidden_size (int): the hidden dimension of the model, importantly this is also the input and output d

Preset configuration for a known MoE model family.

Qwen2-VL-style splice: visual tokens enclosed by vision_start/end tokens. Qwen2-VL wraps each image's visual tokens with a pair of special bounda

Sequence-parallel wrapper for an opaque ViT attention module. Parameters ---------- attn: The original ViT attention layer (any `

Wrapper around the BLAS matmul kernel for FP16/BF16/FP32 for CUDA/RoCM. Performs z = x @ y

CUDA Kernel implementation that will perform rotary position embeddings on the queries and keys before copying into a blocked KV cache.

CUDA Kernel implementation that will perform rotary position embeddings on the queries and keys before copying into a blocked KV cache.

Floating point layer norm kernel for CUDA/RoCM. Performs: z = ln(x)

Sets parameters for DeepSpeed Inference Engine.

Sets parameters for various monitoring methods.

Rollout engine using DeepSpeed hybrid engine. Args: engine: DeepSpeed engine wrapping the model. tokenizer: HuggingFace tokenizer

Configuration settings for LoRAOptimizedLinear. Attributes: lora_r (int): LoRA attention dimension, also known as the rank. Defaults

Optimized version of nn.Linear that adds features such as: * LoRA w. base weight sharding * FP [6,8,12] quantization Arguments:

CUDA Kernel implementation for gather the hidden states of the final token of each sequence. This is used to reduce the cost of the performin

Concrete subclass that appends visual tokens after text tokens.

Configuration class to store the configuration of a `BertModel`.

Config class for DSMoEBase