Implementation:InternLM Lmdeploy UnfusedAttentionKernels

Knowledge Sources	InternLM_Lmdeploy
Domains	GPU_Kernels, Attention
Last Updated	2026-02-07 15:00 GMT

Overview

CUDA kernel declarations for unfused (multi-step) attention operations including masked softmax, QKV bias-transpose, and KV cache transposition.

Description

This header declares kernel functions for the non-fused attention computation path. invokeMaskedSoftmax() applies a masked softmax over attention scores with configurable batch, head, and sequence dimensions. invokeTransposeQKV() and invokeTransposeAttentionOutRemovePadding() handle the transposition of QKV tensors between (batch, seq, heads, dim) and (batch, heads, seq, dim) layouts, with optional padding removal. invokeAddFusedQKVBiasTranspose() splits a fused QKV projection into separate Q, K, V buffers while applying bias and rotary position embedding (RoPE). invokeTranspose4d() and invokeTranspose4dBatchMajor() handle KV cache layout transformations. invokeAddRelativeAttentionBias() and invokeMaskedSoftMaxWithRelPosBias() support relative position encoding schemes. The MaskedSoftmaxParam struct encapsulates softmax parameters.

Usage

Use these kernels in the unfused attention path of transformer inference, where QKV projection, softmax, and output transposition are performed as separate kernel launches rather than a single fused kernel.

Code Reference

Source Location

• Repository: InternLM_Lmdeploy
• File: src/turbomind/kernels/unfused_attention_kernels.h

Signature

template<typename T>
struct MaskedSoftmaxParam {
    T* attention_score; const float* qk; const T* attention_mask;
    int batch_size, q_length, k_length, num_heads;
};

template<typename T>
void invokeMaskedSoftmax(MaskedSoftmaxParam<T>& param, cudaStream_t stream);

template<typename T>
void invokeAddFusedQKVBiasTranspose(
    T* q_buf, T* k_buf, T* v_buf, T* QKV, const T* qkv_bias,
    const int* padding_offset, const int* context_length, const int* input_length,
    const float* rope_theta, const int batch_size, const int seq_len,
    const int token_num, const int head_num, const int kv_head_num,
    const int size_per_head, const int rotary_embedding_dim,
    float rotary_embedding_base, int max_position_embeddings,
    bool use_dynamic_ntk, bool use_logn_attn, cudaStream_t stream);

template<typename T>
void invokeTransposeQKV(T* dst, T* src, const int batch_size, const int seq_len,
    const int head_num, const int size_per_head, const float* scale,
    const int int8_mode, cudaStream_t stream);

Import

#include "src/turbomind/kernels/unfused_attention_kernels.h"

I/O Contract

Inputs

Name	Type	Required	Description
QKV	T*	Yes	Fused QKV projection output
qkv_bias	const T*	No	Bias for Q, K, V projections
attention_mask	const T*	No	Attention mask tensor
rope_theta	const float*	No	RoPE frequency parameters
batch_size, seq_len, head_num, size_per_head	int	Yes	Tensor dimensions

Outputs

Name	Type	Description
q_buf, k_buf, v_buf	T*	Separated and transposed Q, K, V buffers
attention_score	T*	Softmax-normalized attention weights

Usage Examples

using namespace turbomind;

// Split fused QKV with RoPE
invokeAddFusedQKVBiasTranspose(
    q_buf, k_buf, v_buf, qkv_output, qkv_bias,
    padding_offset, context_len, input_len, rope_theta,
    batch_size, seq_len, token_num, n_heads, n_kv_heads,
    head_dim, rotary_dim, rope_base, max_pos, false, false, stream);

// Masked softmax
MaskedSoftmaxParam<half> param{attn_scores, qk, mask, bs, q_len, k_len, n_heads};
invokeMaskedSoftmax(param, stream);

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime