Implementation:Deepspeedai DeepSpeed Inference PT Binding

Knowledge Sources	DeepSpeed
Domains	Inference, PyTorch_Bindings, Transformer, Attention
Last Updated	2026-02-09 00:00 GMT

Overview

PyTorch bindings for optimized transformer inference operations including attention softmax, context computation, and workspace management for autoregressive generation.

Description

This module provides the Python interface to DeepSpeed's inference-optimized CUDA kernels for transformer models. It exposes functions for attention mechanisms (including fused QKV projection, multi-head attention with ALiBi support, and softmax with various masking strategies), workspace allocation for KV-caching, and einsum operations for efficient tensor contractions. The implementation automatically detects model type (GPT vs BERT) from attention mask dimensions and applies appropriate optimizations. The workspace management system dynamically allocates memory based on available GPU memory, supporting both prompt processing and autoregressive token generation with configurable sequence length limits. Special features include triangular causal masking, local attention windowing, layer-specific scaling for model parallelism, and unfused attention fallback for compatibility.

Usage

Use these bindings when deploying transformer models for inference with DeepSpeed's optimizations. The functions handle both prompt processing (long sequences) and generation (incremental decoding with KV-cache), providing significant speedup over standard PyTorch operations through kernel fusion and memory layout optimizations.

Code Reference

Source Location

• Repository: DeepSpeed
• File: csrc/transformer/inference/csrc/pt_binding.cpp

Signature

// Workspace allocation
template <typename T>
void allocate_workspace(unsigned hidden_dim, unsigned num_heads,
                       unsigned prompt_length, unsigned batch_size,
                       unsigned num_layers, unsigned mp_size = 1,
                       bool external_cache = false, unsigned rank = 0,
                       unsigned max_out_tokens = 1024,
                       unsigned min_out_tokens = 1);

// Attention softmax with masking
template <typename T>
at::Tensor ds_softmax(at::Tensor& attn_scores, at::Tensor& attn_mask,
                     at::Tensor& alibi, bool triangular, bool recompute,
                     bool local_attention, int window_size,
                     bool async_op, float layer_scale,
                     int head_offset, int mp_size);

// Multi-head softmax context (full attention)
template <typename T>
std::vector<at::Tensor> ds_softmax_context1(
    at::Tensor& query, at::Tensor& prev_key, at::Tensor& new_key,
    at::Tensor& attn_mask, at::Tensor& prev_value, at::Tensor& new_value,
    int heads, float norm_factor, bool merging, bool triangular,
    bool local_attention, int window_size, bool no_masking);

// Einsum for tensor contractions
template <typename T>
at::Tensor einsum_sec_sm_ecm(at::Tensor& Q, at::Tensor& W);

Import

import deepspeed.ops.transformer.inference as ds_inference

I/O Contract

Parameter	Type	Description
query/key/value	torch.Tensor	Attention tensors [batch×heads×seq×head_dim]
attn_mask	torch.Tensor	Attention mask (2D for BERT, 4D for GPT)
alibi	torch.Tensor	ALiBi position biases (optional)
triangular	bool	Apply causal masking
local_attention	bool	Enable sliding window attention
window_size	int	Local attention window size
layer_scale	float	Per-layer scaling factor (for MP)

Output	Type	Description
attention_output	torch.Tensor	Context vectors [batch×heads×seq×head_dim]
softmax_output	torch.Tensor	Attention probabilities (if requested)

Usage Examples

Workspace Allocation for Generation:

import torch
import deepspeed.ops.transformer.inference.op_binding as ds_ops

# Model configuration
hidden_size = 4096
num_heads = 32
num_layers = 32
batch_size = 8
prompt_length = 512
max_gen_length = 2048

# Allocate workspace for KV-cache and intermediate activations
ds_ops.allocate_workspace_fp16(
    hidden_dim=hidden_size,
    num_heads=num_heads,
    prompt_length=prompt_length,
    batch_size=batch_size,
    num_layers=num_layers,
    mp_size=1,  # Model parallel size
    external_cache=False,  # Manage cache internally
    rank=0,
    max_out_tokens=max_gen_length,
    min_out_tokens=prompt_length
)

print(f"Max sequence length: {ds_ops.get_max_seq_len()}")

GPT-Style Causal Attention:

def gpt_attention_with_cache(query, key, value, past_key, past_value,
                             attention_mask, layer_id, num_heads):
    """
    Compute attention with KV-cache for autoregressive generation
    """
    batch_size = query.size(0)
    seq_length = query.size(1)
    head_dim = query.size(2) // num_heads

    # Reshape for multi-head attention
    q = query.view(batch_size, seq_length, num_heads, head_dim)
    k = key.view(batch_size, seq_length, num_heads, head_dim)
    v = value.view(batch_size, seq_length, num_heads, head_dim)

    # Concatenate with past
    full_key = torch.cat([past_key, k], dim=1) if past_key is not None else k
    full_value = torch.cat([past_value, v], dim=1) if past_value is not None else v

    # Compute attention with fused softmax
    context, key_cache, value_cache = ds_ops.softmax_context_fp16(
        query=q,
        prev_key=full_key if past_key is not None else torch.empty(0),
        new_key=k,
        attn_mask=attention_mask,
        prev_value=full_value if past_value is not None else torch.empty(0),
        new_value=v,
        heads=num_heads,
        norm_factor=1.0 / (head_dim ** 0.5),
        merging=past_key is not None,  # Merging new with past
        triangular=True,  # Causal masking
        local_attention=False,
        window_size=0,
        no_masking=False
    )

    return context.view(batch_size, seq_length, -1), key_cache, value_cache

Attention Softmax with ALiBi:

def compute_attention_scores_alibi(scores, alibi_bias, layer_id,
                                  num_heads, use_causal_mask=True):
    """
    Apply softmax with ALiBi position biases
    """
    batch, heads, seq_q, seq_k = scores.shape

    # Create causal mask if needed
    if use_causal_mask:
        causal_mask = torch.triu(
            torch.ones(seq_q, seq_k, device=scores.device) * float('-inf'),
            diagonal=1
        ).unsqueeze(0).unsqueeze(0)
    else:
        causal_mask = torch.zeros(1, 1, 1, 1, device=scores.device)

    # Apply softmax with ALiBi and masking
    attn_probs = ds_ops.softmax_fp16(
        attn_scores=scores,
        attn_mask=causal_mask,
        alibi=alibi_bias,
        triangular=use_causal_mask,
        recompute=False,
        local_attention=False,
        window_size=0,
        async_op=False,
        layer_scale=1.0,
        head_offset=layer_id * num_heads,  # For model parallelism
        mp_size=1
    )

    return attn_probs

Local/Sliding Window Attention:

def sliding_window_attention(query, key, value, window_size=256):
    """
    Efficient local attention for long sequences
    """
    batch, seq_len, hidden = query.shape
    num_heads = 16
    head_dim = hidden // num_heads

    q = query.view(batch, seq_len, num_heads, head_dim)
    k = key.view(batch, seq_len, num_heads, head_dim)
    v = value.view(batch, seq_len, num_heads, head_dim)

    # Compute attention scores
    scores = torch.matmul(q.transpose(1, 2), k.transpose(1, 2).transpose(-2, -1))
    scores = scores / (head_dim ** 0.5)

    # Apply local attention masking (only attend to window_size tokens)
    mask = torch.zeros(1, 1, seq_len, seq_len, device=query.device)
    for i in range(seq_len):
        start = max(0, i - window_size)
        mask[0, 0, i, :start] = float('-inf')

    # Fused softmax with local attention
    attn_probs = ds_ops.softmax_fp16(
        attn_scores=scores,
        attn_mask=mask,
        alibi=torch.empty(0),
        triangular=False,
        recompute=False,
        local_attention=True,
        window_size=window_size,
        async_op=False,
        layer_scale=1.0,
        head_offset=0,
        mp_size=1
    )

    # Apply to values
    context = torch.matmul(attn_probs, v.transpose(1, 2))
    return context.transpose(1, 2).contiguous().view(batch, seq_len, hidden)

BERT-Style Bidirectional Attention:

def bert_attention(query, key, value, attention_mask, num_heads):
    """
    BERT bidirectional attention (no causal masking)
    """
    batch, seq_len, hidden = query.shape
    head_dim = hidden // num_heads

    # Reshape for multi-head
    q = query.view(batch, seq_len, num_heads, head_dim)
    k = key.view(batch, seq_len, num_heads, head_dim)
    v = value.view(batch, seq_len, num_heads, head_dim)

    # BERT mask is 2D: [batch, seq_len]
    # Convert to attention mask format
    extended_mask = attention_mask.unsqueeze(1).unsqueeze(2)
    extended_mask = (1.0 - extended_mask) * -10000.0

    # Compute attention
    context, _, _ = ds_ops.softmax_context_fp16(
        query=q,
        prev_key=torch.empty(0),
        new_key=k,
        attn_mask=extended_mask,
        prev_value=torch.empty(0),
        new_value=v,
        heads=num_heads,
        norm_factor=1.0 / (head_dim ** 0.5),
        merging=False,
        triangular=False,  # Bidirectional
        local_attention=False,
        window_size=0,
        no_masking=False
    )

    return context.view(batch, seq_len, hidden)

Related Pages

• Inference Context - Workspace and state management
• Inference cuBLAS Wrappers - GEMM operations used internally
• Custom CUDA Layers - Underlying kernel implementations