Our first walk of microsoft/DeepSpeed stayed on the training side of the tree: ZeRO’s three sharding stages, the pipe engine’s micro-batch schedules, the launcher that SSHes into clusters, and the accelerator layer that speaks for six hardware vendors. That tour deliberately skipped an entire second codebase hiding under the same roof: deepspeed/inference, and inside it a ground-up rebuild called Inference v2. Serving a trained model is a different computational shape from training it — instead of one big batch of identical-length sequences, a server juggles dozens of half-finished conversations, each with its own prompt length, its own growing KV history, and its own arrival time. Training’s padded tensors waste both memory and arithmetic in that world. The v2 code answers with a data layout it calls ragged: many variable-length sequences packed edge to edge into single tensors, with per-sequence metadata travelling alongside in parallel buffers. As always, this is an educational read of published source on the master branch, Apache 2.0 licensed; the training-side tour covers the modules we now zoom past, and readers comparing serving designs can also see how vLLM and SGLang attack the same problem.

The one-sentence summary: InferenceEngineV2 is the front desk that admits requests, DSStateManager is the ledger that tracks every sequence and every KV block, RaggedBatchWrapper is the crate that packs admitted requests into one GPU-shaped batch, and a registry of swappable DSModules wired to hand-written CUDA kernels does the arithmetic.

Overview of the v2 serving stack: a factory pairs a checkpoint source with a model policy, the engine admits requests through the state manager, ragged batching carries sequences into the Llama model, whose layers are assembled by module heuristics and executed by blocked CUDA kernels.

Reading the overview from left to right:

• build_hf_engine is the public constructor, one call from a Hugging Face model name to a running engine.
• HuggingFaceCheckpointEngine streams weights from the hub, preferring safetensors, with the file fetch at line 38.
• InferenceV2Policy is the abstract strategy that knows a model family: how to build it and where every parameter belongs.
• InferenceEngineV2 exposes put at line 107, query at line 158, and can_schedule at line 184.
• DSStateManager creates sequences with get_or_create_sequence at line 130 and hands out storage with allocate_blocks at line 203.
• DSSequenceDescriptor keeps the block list per request, extended by extend_kv_cache at line 235.
• BlockedKVCache owns the pools, with offload and restore moving blocks to host memory.
• RaggedBatchWrapper turns sequence descriptors into device buffers via insert_sequence at line 134 and finalize at line 184.
• Llama2InferenceModel walks embed, transformer layers, and unembed over the ragged batch.
• DSTransformerModelBase builds every layer through factory methods like make_qkv_layer at line 363.

The front desk: put, query, can_schedule

Three verbs define the engine’s public surface, and each maps to a distinct phase of a request’s life. put is admission: the caller supplies a unique id and a tensor of prompt tokens, and the engine first stages them into a host-side placeholder — PlaceholderSequenceDescriptor records what is coming without touching device memory — then asks the state manager for a real DSSequenceDescriptor and copies tokens into a pinned host buffer so the next forward pass can sweep them up. query is the opposite direction: given a uid and a budget of new tokens and blocks, it reports whether the sequence can run this step and hands back the token slice to feed. The scheduler’s judgment call lives in can_schedule, which walks the candidate uids against two limits — the tracked-sequence cap and the free-block count — and returns a SchedulingResult instead of throwing, so a caller can back off gracefully when the engine is saturated. flush retires a finished conversation, returning its blocks through flush_sequence, and serialize saves the whole warm engine for later reload through build_engine_from_ds_checkpoint.

The ledger underneath is DSStateManager, constructed from a DSStateManagerConfig that fixes how many sequences can be tracked and how many cache groups exist. Memory accounting is deliberately indirect: nothing here stores a KV tensor per sequence. Instead each descriptor remembers integer block ids, and the pools themselves live in BlockedKVCache, one per cache group, shaped as flat block-indexed tensors. Handing out ids is BlockedAllocator, a bookkeeping class whose allocate at line 50 returns a tensor of free indices and whose free at line 74 takes them back — the allocator is the whole memory manager, and its discipline is what makes overflow errors impossible to hide. The cache adds two moves the allocator does not know about: offload copies chosen blocks to host memory to free device room, and restore brings them back when their sequence is rescheduled, with a shape helper called split_kv separating key from value views along the way.

Ragged batches: packing a whole server into four tensors

The batch format is where the design earns its name. A padded batch pads every request to the longest one and computes the padding too; a ragged batch concatenates the actual tokens of all in-flight sequences into a single one-dimensional stream and carries, beside it, a metadata buffer saying where each sequence’s tokens start and stop. RaggedBatchWrapper builds exactly this: insert_sequence appends one descriptor’s tokens and block pointers into preallocated host buffers, and finalize uploads everything with a handful of non-blocking copies, after which the batch exposes parallel device tensors — input_ids, the metadata buffer, tokens-to-sequence maps, and crucially kv_ptrs, a flat tensor of raw pointers into the block pools, one entry per layer and per sequence. That pointer table is the hinge of the whole design: the CUDA kernels never see a Python-side index, they read block ids from the metadata buffer and address the KV pools directly, which is why the same kernel handles any batch composition without recompilation.

Attention then proceeds in blocks. AtomBuilder derives per-head work atoms — fixed-size query-and-key spans — from the batch metadata, and the module calls BlockedFlashAttn with those atoms, out scratch, and the pointer tables; its call signature at line 88 is worth reading as the specification of the whole ragged layout. The sizes are tuned by get_q_block_size per head dimension. The same pattern frames both ends of the network: RaggedEmbeddingKernel gathers embedding rows for the packed token stream, and RaggedLogitsGather copies only the last-token rows of each sequence out for sampling — the unembed GEMM still runs over every hidden state, but the gather that produces logits touches the minimum.

Policies, containers, and a metaclass that checks your weights

Getting a Llama checkpoint into this layout is the job of the policy layer. InferenceV2Policy declares two responsibilities — build_model at line 140 instantiates the model class, and populate_model_parameters at line 189 routes every incoming tensor to the right destination through a ContainerMap that separates transformer-layer containers from the embedding-side one. A metaclass enforces the contract: LayerMetaclass rewrites class construction so that annotated attributes become parameter slots, and its call override at line 150 refuses half-initialized containers — a checkpoint loader cannot silently leave a layer empty. Weights themselves are wrapped as InferenceParameter, a tensor subclass that tracks device residency so a quantized weight still living on CPU is caught at use time rather than at a mysterious kernel crash.

The concrete pair for this tour is the Llama family. Llama2Policy supplies the model class and its container map; Llama2TransformerContainer and Llama2NonTransformerContainer name the per-layer and global weights. The model’s forward at line 199 is refreshingly literal: _forward_embed, then a Python loop calling _forward_transformer_layer once per layer index, then _forward_unembed. All the cleverness is one level down, in DSTransformerModelBase, whose factory methods — make_embedding_layer at line 303, make_qkv_layer at line 363, make_norm_layer at line 579 — consult per-architecture configuration, and whose maybe_allocate_kv at line 436 asks the state manager for blocks before each pass using get_kv_requirements. A sibling base, DSMoETransformerModelBase, adds expert layers through make_moe_layer for Mixtral and Qwen-MoE checkpoints.

DSModules: every organ is swappable

The most transferable idea in the tree is the module contract. DSModuleBase requires each implementation to name itself and answer supports_config — a predicate saying which configurations it can handle. Registration is self-service: subclassing DSModuleRegistryBase enrolls the class via register_module at line 49, so importing an implementation file is enough to make it discoverable. At engine construction, the heuristic pickers — instantiate_attention at line 36, instantiate_linear at line 75, instantiate_pre_norm at line 158, plus their embed, unembed, and MoE siblings — walk the registered implementations in priority order and take the first that supports the requested configuration. The shipped set shows the spectrum: DSDenseBlockedAttention builds atoms with build_atoms at line 139 before its forward at line 152; BlasFPLinear is the honest pure-GEMM fallback for unquantized weights; DSPreRMSCUDAModule fuses RMSNorm; DSRaggedEmbedding and DSRaggedUnembed bracket the stream. Anyone adding a model family follows a documented recipe — the AddingAModel guide in the same folder — and anyone adding a faster kernel only has to satisfy the same config predicate, which is exactly how this tree wants extensions to arrive.

From Install to Your First Model

Install with pip install deepspeed, then the v2 path is three lines of Python: from deepspeed.inference.v2.engine_factory import build_hf_engine; engine = build_hf_engine("meta-llama/Llama-2-7b-hf"); engine.put(0, tokens). Between those lines, the factory loads the checkpoint streaming class, builds the policy, populates every container, wraps the model in InferenceEngineV2, and returns an object whose put/query/flush verbs are the entire API. Be honest about the fine print, because the code is: the v2 engine is marked experimental in its own docs, the blocked CUDA kernels target NVIDIA tensor-core GPUs, and running with tensor parallelism greater than one is the configuration the engine actually exercises — single-GPU mode works through the BLAS linear module but is not the optimized path. Quantized-weight support exists in the cutlass GEMM kernels but the coverage of bit-widths is narrower than the training side’s optimizer zoo, and there is no continuous-batching scheduler loop shipped as a runnable daemon — the engine provides put/query/can_schedule and expects your serving loop to call them, unlike the turnkey HTTP servers this series has met in lmdeploy and local inference servers. What you get in exchange is the clearest executable specification of ragged serving we have read: four trees of code where every tensor has a declared owner. Compare the training half at our first DeepSpeed tour, the quantized-kernel approach of Triton, the execution-plan machinery of ONNX Runtime, the mobile side of ExecuTorch, and the speech sibling faster-whisper — the same lesson repeats: production inference is bookkeeping with arithmetic inside.

Enjoyed this post? Never miss out on future posts by following us