Implementation:Deepspeedai DeepSpeed Inference Context

Knowledge Sources	DeepSpeed
Domains	Inference, Memory_Management, Context_Management, CUDA
Last Updated	2026-02-09 00:00 GMT

Overview

Singleton context manager for inference operations providing workspace allocation, cuBLAS handle management, stream coordination, and KV-cache memory management.

Description

The InferenceContext class implements a singleton pattern that manages global state for DeepSpeed inference operations. It provides centralized workspace allocation with automatic sizing based on available GPU memory, model configuration (hidden size, number of heads/layers), and sequence length requirements. The workspace accommodates attention computations, KV-cache storage (optionally external), and intermediate activations with dynamic adjustment for prompt processing versus token generation. The class maintains a cuBLAS handle configured with tensor core math mode, manages CUDA streams for compute and communication operations with event-based synchronization, and tracks generation state including current token count and seed for random number generation. Memory allocation is lazy and can be released/retaken for flexible resource management in deployment scenarios.

Usage

Access the singleton instance when implementing inference kernels that need workspace memory, cuBLAS operations, or stream management. The context automatically handles memory constraints and provides diagnostic output for workspace sizing decisions, crucial for deploying large models efficiently.

Code Reference

Source Location

• Repository: DeepSpeed
• File: csrc/transformer/inference/includes/inference_context.h

Signature

class InferenceContext {
public:
    static InferenceContext& Instance();

    // Workspace management
    void GenWorkSpace(const unsigned& num_layers,
                     const unsigned& num_heads,
                     const size_t& batch_size,
                     const size_t& prompt_len,
                     const size_t& hidden_dim,
                     const unsigned& mp_size,
                     const bool& external_cache,
                     const size_t& elem_size,
                     const unsigned& rank,
                     unsigned max_out_tokens,
                     unsigned min_out_tokens);

    void* GetWorkSpace();
    void* GetAttentionUnfusedWorkspace();
    size_t get_workspace_size() const;
    size_t GetMaxTokenLength() const;

    // cuBLAS management
    cublasHandle_t GetCublasHandle();

    // Stream management
    cudaStream_t GetCurrentStream(bool other_stream = false);
    cudaStream_t GetCommStream(bool async_op = false);
    void SynchComp();  // Synchronize compute and comm streams
    void SynchComm();

    // Generation state
    unsigned current_tokens() const;
    void advance_tokens();
    void reset_tokens(unsigned initial_tokens = 1);
    std::pair<uint64_t, uint64_t> IncrementOffset(uint64_t offset_inc);
    void SetSeed(uint64_t new_seed);

    // Workspace control
    void release_workspace();
    bool retake_workspace();
};

Import

#include "inference_context.h"

I/O Contract

Parameter	Type	Description
num_layers	unsigned	Number of transformer layers
num_heads	unsigned	Number of attention heads
batch_size	size_t	Batch size for inference
prompt_len	size_t	Initial prompt length
hidden_dim	size_t	Model hidden dimension
mp_size	unsigned	Model parallel size
max_out_tokens	unsigned	Maximum total sequence length

Output	Type	Description
workspace	void*	Allocated workspace pointer
cublas_handle	cublasHandle_t	Configured cuBLAS handle
streams	cudaStream_t	CUDA streams for operations

Usage Examples

Basic Context Initialization:

#include "inference_context.h"

// Model configuration
unsigned num_layers = 32;
unsigned num_heads = 32;
size_t hidden_dim = 4096;
size_t batch_size = 8;
size_t prompt_len = 512;
unsigned max_tokens = 2048;

// Initialize workspace
InferenceContext::Instance().GenWorkSpace(
    num_layers, num_heads, batch_size, prompt_len, hidden_dim,
    1,      // mp_size (no model parallelism)
    false,  // external_cache (manage internally)
    sizeof(__half),  // element size
    0,      // rank
    max_tokens,      // maximum sequence length
    prompt_len       // minimum sequence length
);

// Get workspace for kernel
void* workspace = InferenceContext::Instance().GetWorkSpace();
size_t workspace_size = InferenceContext::Instance().get_workspace_size();
printf("Allocated workspace: %.2f GB\n",
       (float)workspace_size / (1024*1024*1024));

Using cuBLAS Handle:

// Perform matrix multiplication using context's cuBLAS handle
void inference_matmul(__half* A, __half* B, __half* C,
                     int M, int N, int K) {
    auto handle = InferenceContext::Instance().GetCublasHandle();
    auto stream = InferenceContext::Instance().GetCurrentStream();

    cublasSetStream(handle, stream);

    float alpha = 1.0f, beta = 0.0f;
    cublas_gemm_ex(handle, CUBLAS_OP_N, CUBLAS_OP_N,
                   M, N, K, &alpha, &beta, A, B, C,
                   CUBLAS_GEMM_DEFAULT_TENSOR_OP);
}

Stream Coordination:

// Overlap computation with communication
void pipeline_forward(__half* input, __half* output,
                     __half* weights, int size) {
    auto comp_stream = InferenceContext::Instance().GetCurrentStream();
    auto comm_stream = InferenceContext::Instance().GetCommStream(true);

    // Launch computation on compute stream
    compute_kernel<<<blocks, threads, 0, comp_stream>>>(
        input, output, weights, size);

    // Synchronize compute before communication
    InferenceContext::Instance().SynchComp();

    // Launch communication (all-gather/all-reduce)
    communication_kernel<<<1, 1, 0, comm_stream>>>(output, size);

    // Wait for communication before next compute
    InferenceContext::Instance().SynchComm();
}

Token Generation Tracking:

class AutoregressiveGenerator {
    InferenceContext& ctx = InferenceContext::Instance();

public:
    void generate(int num_tokens_to_generate) {
        ctx.reset_tokens(1);  // Start with 1 token (BOS)

        for (int i = 0; i < num_tokens_to_generate; i++) {
            int current = ctx.current_tokens();
            printf("Generating token %d\n", current);

            // Forward pass generates next token
            forward_step();

            // Advance token counter
            ctx.advance_tokens();

            // Check if we've reached max sequence length
            if (ctx.current_tokens() > ctx.GetMaxTokenLength()) {
                fprintf(stderr, "Exceeded max sequence length!\n");
                break;
            }
        }
    }
};

Dynamic Workspace Allocation:

// Adjust workspace for different batch sizes
void reconfigure_workspace(size_t new_batch_size,
                          size_t new_seq_len) {
    auto& ctx = InferenceContext::Instance();

    // Release current workspace
    ctx.release_workspace();

    // Reconfigure with new parameters
    ctx.GenWorkSpace(
        32,  // num_layers
        32,  // num_heads
        new_batch_size,
        new_seq_len,  // prompt_len
        4096,  // hidden_dim
        1,  // mp_size
        false,  // external_cache
        sizeof(__half),
        0,  // rank
        2048,  // max_out_tokens
        new_seq_len  // min_out_tokens
    );

    // Verify workspace was reallocated
    if (!ctx.retake_workspace()) {
        throw std::runtime_error("Failed to reallocate workspace");
    }
}

Attention Workspace Management:

// Use separate workspace regions for different operations
__global__ void attention_kernel(__half* Q, __half* K, __half* V,
                                __half* output, int batch, int heads,
                                int seq_len, int head_dim) {
    // Main workspace for intermediate results
    __half* main_workspace = static_cast<__half*>(
        InferenceContext::Instance().GetWorkSpace());

    // Separate region for unfused attention operations
    __half* attn_workspace = static_cast<__half*>(
        InferenceContext::Instance().GetAttentionUnfusedWorkspace());

    // Use attn_workspace for attention scores
    __half* scores = attn_workspace;
    int scores_size = batch * heads * seq_len * seq_len;

    // Compute attention...
}

External KV-Cache Management:

// When KV-cache is managed externally (e.g., PagedAttention)
class ExternalCacheManager {
    void* kv_cache;
    size_t cache_size;

public:
    ExternalCacheManager(size_t batch, size_t max_seq,
                        size_t num_layers, size_t hidden) {
        // Allocate KV-cache separately
        cache_size = 2 * num_layers * batch * max_seq * hidden * sizeof(__half);
        cudaMalloc(&kv_cache, cache_size);

        // Initialize context with external_cache=true
        InferenceContext::Instance().GenWorkSpace(
            num_layers, hidden / 64, batch, 1, hidden,
            1, true,  // external_cache = true
            sizeof(__half), 0, max_seq, 1
        );
    }

    ~ExternalCacheManager() {
        cudaFree(kv_cache);
    }

    void* get_kv_cache() { return kv_cache; }
};

Memory-Constrained Deployment:

// Handle OOM gracefully
bool try_allocate_workspace(size_t batch, size_t seq_len,
                           size_t max_tokens) {
    try {
        InferenceContext::Instance().GenWorkSpace(
            32, 32, batch, seq_len, 4096,
            1, false, sizeof(__half), 0,
            max_tokens, seq_len
        );
        return true;
    } catch (const std::runtime_error& e) {
        // Workspace allocation failed (OOM)
        fprintf(stderr, "Workspace allocation failed: %s\n", e.what());

        // Try with smaller max_tokens
        if (max_tokens > seq_len * 2) {
            return try_allocate_workspace(batch, seq_len, max_tokens / 2);
        }
        return false;
    }
}

// Usage
if (!try_allocate_workspace(8, 512, 4096)) {
    fprintf(stderr, "Unable to allocate sufficient workspace\n");
    // Fall back to smaller batch or sequence length
}

Random Number Generation:

// Coordinate RNG state across operations
void set_generation_seed(uint64_t seed) {
    InferenceContext::Instance().SetSeed(seed);
}

__global__ void dropout_kernel(__half* data, uint8_t* mask,
                               int size, float ratio) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= size) return;

    // Get unique seed/offset for this operation
    auto [seed, offset] = InferenceContext::Instance().IncrementOffset(size);

    // Use seed/offset for curand initialization
    curandState state;
    curand_init(seed, idx, offset, &state);

    float rand_val = curand_uniform(&state);
    mask[idx] = (rand_val > ratio) ? 1 : 0;
    data[idx] = mask[idx] ? data[idx] / (1.0f - ratio) : 0.0f;
}

Related Pages

• Inference PT Binding - Uses context for workspace
• Inference cuBLAS Wrappers - cuBLAS operations via context
• GEMM Test - Algorithm tuning with context handle