Implementation:Deepspeedai DeepSpeed Custom CUDA Layers

Knowledge Sources	DeepSpeed
Domains	Transformer, CUDA_Kernels, Neural_Networks, Training
Last Updated	2026-02-09 00:00 GMT

Overview

CUDA kernel launch declarations for fused transformer operations including attention, layer normalization, GELU activation, and dropout.

Description

This header defines the launch interfaces for custom CUDA kernels implementing optimized transformer building blocks. It provides fused operations that combine multiple steps (e.g., bias addition with activation functions, residual connections with layer normalization) to reduce memory bandwidth and kernel launch overhead. The implementation supports both forward and backward passes for training, with specialized kernels for attention mechanisms including softmax with masking, QKV transformations, and context computation. The header also includes utilities for dynamic token selection and mask manipulation used in sparse attention patterns and length-adaptive training.

Usage

Use these kernel launch functions when implementing high-performance transformer layers in CUDA. The fused operations provide significant speedup compared to separate kernel launches for each operation. Particularly useful for training large language models where attention and feed-forward operations dominate computation time.

Code Reference

Source Location

• Repository: DeepSpeed
• File: csrc/includes/custom_cuda_layers.h

Signature

// Fused bias + GELU activation
template <typename T>
void launch_bias_gelu(const T* input, const T* bias, T* output,
                     int intermediate_size, int batch_size, cudaStream_t stream);

// Fused bias + residual + layer norm
template <typename T>
void launch_bias_residual_layer_norm(T* vals, const T* residual,
                                    const T* gamma, const T* beta,
                                    float epsilon, int batch_size,
                                    int hidden_dim, cudaStream_t stream,
                                    bool preLayerNorm, bool training,
                                    T* vars, T* means);

// Attention softmax with mask and scaling
template <typename T>
void launch_attn_softmax(T* vals, const T* attn_mask,
                        int batch_size, int heads, int sequence_length,
                        cudaStream_t stream);

// QKV transformation: [0,1,2,3] -> [0,2,1,3]
template <typename T>
void launch_bias_add_transform_0213(T* outputs, const T* vals, const T* bias,
                                   int batch_size, int seq_length,
                                   int hidden_dim, int heads,
                                   cudaStream_t stream, int trans_count);

// Dropout with mask generation
template <typename T>
void launch_dropout(T* vals, const T* bias, uint8_t* mask,
                   int batch, int dim, float ratio, cudaStream_t stream);

Import

#include "csrc/includes/custom_cuda_layers.h"

I/O Contract

Parameter	Type	Description
input/vals	T*	Input tensor data (float/__half)
bias	const T*	Bias vector to add
output	T*	Output tensor buffer
batch_size	int	Number of sequences
hidden_dim	int	Hidden dimension size
heads	int	Number of attention heads
sequence_length	int	Sequence length
stream	cudaStream_t	CUDA stream for async execution

Output	Type	Description
output	T*	Transformed/computed results
mask	uint8_t*	Dropout mask (if applicable)

Usage Examples

Fused Bias + GELU:

// Feed-forward network intermediate layer
int batch_size = 32;
int seq_len = 512;
int intermediate_size = 4096;
int tokens = batch_size * seq_len;

__half *ff_output, *bias, *gelu_output;
// ... allocate memory ...

launch_bias_gelu(ff_output, bias, gelu_output,
                intermediate_size, tokens, stream);

Attention with Mask:

// Apply softmax to attention scores with causal mask
int batch_size = 8;
int num_heads = 16;
int seq_length = 1024;

__half *attn_scores, *attn_mask;
// ... compute attention scores ...

launch_attn_softmax(attn_scores, attn_mask,
                   batch_size, num_heads, seq_length, stream);

Bias + Residual + LayerNorm:

// Post-attention layer normalization with residual connection
__half *attn_output, *residual, *gamma, *beta, *normalized;
__half *variance, *mean;
float epsilon = 1e-5;
int batch_size = 32;
int hidden_dim = 1024;

launch_bias_residual_layer_norm(attn_output, residual,
                                gamma, beta, epsilon,
                                batch_size, hidden_dim, stream,
                                true, true, variance, mean);

QKV Transform for Multi-Head Attention:

// Reshape and transpose for multi-head attention
__half *qkv_combined, *bias, *q_transformed;
int batch = 16, seq_len = 512, hidden = 768, heads = 12;

launch_bias_add_transform_0213(q_transformed, qkv_combined, bias,
                               batch, seq_len, hidden, heads,
                               stream, 3);  // 3 for Q, K, V

Related Pages

• Transformer CUDA - Uses these kernels for transformer implementation
• Normalize Layer - C++ wrapper for layer norm kernels
• Reduction Utils - Reduction primitives used in these kernels