Implementation:Deepspeedai DeepSpeed Inference Cublas Wrappers

Knowledge Sources	DeepSpeed
Domains	Linear_Algebra, Inference, CUDA_Kernels, Performance_Optimization
Last Updated	2026-02-09 00:00 GMT

Overview

Cross-platform cuBLAS/rocBLAS wrappers providing unified GEMM interfaces for FP32, FP16, and BF16 matrix multiplication with algorithm selection.

Description

This header provides a unified API layer abstracting differences between NVIDIA cuBLAS and AMD rocBLAS libraries, enabling portable GEMM operations across CUDA and ROCm platforms. It implements wrapper functions for standard GEMM (cublas_gemm_ex) and strided batched GEMM (cublas_strided_batched_gemm) with support for mixed-precision computation (FP16/BF16 inputs with FP32 accumulation). The implementation uses conditional compilation to handle API differences between platforms, including version-specific compatibility for PyTorch 2.0 transitions on ROCm. Algorithm selection is exposed as a parameter allowing performance tuning based on problem dimensions, and the wrappers include error checking with detailed diagnostics for matrix dimension debugging.

Usage

Use these wrappers when implementing inference kernels that need to run on both NVIDIA and AMD GPUs, or when you need explicit control over cuBLAS algorithm selection for performance optimization. The strided batched variant is particularly important for transformer attention mechanisms where multiple matrix multiplications with regular stride patterns are computed in parallel.

Code Reference

Source Location

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

Signature

// FP32 GEMM: C = alpha * op(A) * op(B) + beta * C
int cublas_gemm_ex(cublasHandle_t handle,
                   cublasOperation_t transa, cublasOperation_t transb,
                   int m, int n, int k,
                   const float* alpha, const float* beta,
                   const float* A, const float* B, float* C,
                   cublasGemmAlgo_t algo, int b_stride = -1);

// FP16/BF16 GEMM with FP32 accumulation
template <typename T>
int cublas_gemm_ex(cublasHandle_t handle,
                   cublasOperation_t transa, cublasOperation_t transb,
                   int m, int n, int k,
                   const float* alpha, const float* beta,
                   const T* A, const T* B, T* C,
                   cublasGemmAlgo_t algo, int b_stride = -1);

// Strided batched GEMM for multiple independent matrices
template <typename T>
int cublas_strided_batched_gemm(cublasHandle_t handle,
                                int m, int n, int k,
                                const float* alpha, const float* beta,
                                const T* A, const T* B, T* C,
                                cublasOperation_t op_A, cublasOperation_t op_B,
                                int stride_A, int stride_B, int stride_C,
                                int batch, cublasGemmAlgo_t algo);

Import

#include "inference_cublas_wrappers.h"

I/O Contract

Parameter	Type	Description
handle	cublasHandle_t	cuBLAS context handle
transa/transb	cublasOperation_t	Transpose flags (N/T)
m, n, k	int	Matrix dimensions (m×k) × (k×n) = (m×n)
alpha, beta	const float*	Scaling factors
A, B	const T*	Input matrices (FP32/FP16/BF16)
C	T*	Output matrix (in-place update)
algo	cublasGemmAlgo_t	Algorithm ID for cuBLAS
stride_A/B/C	int	Memory stride between batches
batch	int	Number of matrix multiplications

Output	Type	Description
C	T*	Result matrix C = alpha*op(A)*op(B) + beta*C
return	int	0 on success, EXIT_FAILURE on error

Usage Examples

Basic FP16 Matrix Multiplication:

#include "inference_cublas_wrappers.h"

cublasHandle_t handle;
cublasCreate(&handle);

int M = 1024, N = 4096, K = 768;
__half *A, *B, *C;
// ... allocate and initialize ...

float alpha = 1.0f, beta = 0.0f;

int status = cublas_gemm_ex(
    handle,
    CUBLAS_OP_N, CUBLAS_OP_N,  // No transpose
    M, N, K,
    &alpha, &beta,
    A, B, C,
    CUBLAS_GEMM_DEFAULT_TENSOR_OP);

if (status != 0) {
    fprintf(stderr, "GEMM failed\n");
}

Attention Score Computation (Strided Batch):

// Compute attention scores: S = Q * K^T for all heads
// Q: [batch*heads, seq_len, head_dim]
// K: [batch*heads, seq_len, head_dim]
// S: [batch*heads, seq_len, seq_len]

int batch_heads = batch_size * num_heads;
int seq_len = 512;
int head_dim = 64;

__half *Q, *K, *scores;
// ... allocate memory ...

float scale = 1.0f / sqrtf((float)head_dim);
float beta = 0.0f;

cublasSetStream(handle, stream);
cublas_strided_batched_gemm(
    handle,
    seq_len,        // M: output rows
    seq_len,        // N: output cols
    head_dim,       // K: reduction dimension
    &scale, &beta,
    K,              // A: keys
    Q,              // B: queries
    scores,         // C: output scores
    CUBLAS_OP_T,    // Transpose keys
    CUBLAS_OP_N,    // Don't transpose queries
    seq_len * head_dim,    // stride_A
    seq_len * head_dim,    // stride_B
    seq_len * seq_len,     // stride_C
    batch_heads,           // batch count
    CUBLAS_GEMM_DEFAULT_TENSOR_OP);

Attention Context Computation:

// Compute context: C = softmax(S) * V
// attn_probs: [batch*heads, seq_len, seq_len]
// V: [batch*heads, seq_len, head_dim]
// context: [batch*heads, seq_len, head_dim]

__half *attn_probs, *V, *context;

float alpha = 1.0f, beta = 0.0f;

cublas_strided_batched_gemm(
    handle,
    head_dim,       // M: head dimension
    seq_len,        // N: sequence length
    seq_len,        // K: attention dimension
    &alpha, &beta,
    V,              // A: values
    attn_probs,     // B: attention probabilities
    context,        // C: output context
    CUBLAS_OP_N,    // Don't transpose V
    CUBLAS_OP_N,    // Don't transpose attn_probs
    seq_len * head_dim,    // stride_A
    seq_len * seq_len,     // stride_B
    seq_len * head_dim,    // stride_C
    batch_heads,
    CUBLAS_GEMM_DEFAULT_TENSOR_OP);

Linear Layer with Algorithm Selection:

// Optimized linear layer with tuned algorithm
class OptimizedLinear {
    cublasHandle_t handle;
    cublasGemmAlgo_t best_algo;
    __half *weight;  // [out_features, in_features]
    int in_features, out_features;

public:
    OptimizedLinear(cublasHandle_t h, int in_f, int out_f,
                   cublasGemmAlgo_t algo)
        : handle(h), in_features(in_f), out_features(out_f),
          best_algo(algo) {
        cudaMalloc(&weight, sizeof(__half) * in_f * out_f);
    }

    void forward(__half* input, __half* output, int batch_size) {
        // Y = X * W^T
        float alpha = 1.0f, beta = 0.0f;

        cublas_gemm_ex(
            handle,
            CUBLAS_OP_T,    // Transpose weight
            CUBLAS_OP_N,    // Don't transpose input
            out_features,   // M
            batch_size,     // N
            in_features,    // K
            &alpha, &beta,
            weight, input, output,
            best_algo);
    }
};

Cross-Platform Inference:

// Platform-agnostic GEMM wrapper
template <typename T>
void inference_matmul(cublasHandle_t handle,
                     const T* A, const T* B, T* C,
                     int m, int n, int k,
                     bool transpose_a, bool transpose_b) {
    float alpha = 1.0f, beta = 0.0f;

    auto op_a = transpose_a ? CUBLAS_OP_T : CUBLAS_OP_N;
    auto op_b = transpose_b ? CUBLAS_OP_T : CUBLAS_OP_N;

    #if defined(__HIP_PLATFORM_AMD__)
    // ROCm-specific algorithm
    auto algo = rocblas_gemm_algo_standard;
    #else
    // NVIDIA-specific algorithm
    auto algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP;
    #endif

    int status = cublas_gemm_ex(
        handle, op_a, op_b, m, n, k,
        &alpha, &beta, A, B, C, algo);

    if (status != 0) {
        throw std::runtime_error("GEMM failed");
    }
}

BFloat16 Support:

#ifdef BF16_AVAILABLE
// Mixed precision: BF16 input/output, FP32 accumulation
void bf16_linear_layer(cublasHandle_t handle,
                      const __nv_bfloat16* input,
                      const __nv_bfloat16* weight,
                      __nv_bfloat16* output,
                      int batch, int in_dim, int out_dim) {
    float alpha = 1.0f, beta = 0.0f;

    cublas_gemm_ex(
        handle,
        CUBLAS_OP_T, CUBLAS_OP_N,
        out_dim, batch, in_dim,
        &alpha, &beta,
        weight, input, output,
        CUBLAS_GEMM_DEFAULT_TENSOR_OP);
}
#endif

Related Pages

• GEMM Test - Algorithm selection and tuning utilities
• Inference PT Binding - Uses these wrappers for attention
• Inference Context - Manages cuBLAS handles