add cutlass fmha support in PackedAttention (#15838)

### Description
<!-- Describe your changes. -->
Support cutlass fMHA in PackedAttention. Though we have fMHA trt kernel,
it doesn't support relative bias position. Cutlass fmha has support for
RBP and also support lower end GPUs(5.3, 6.x).


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
Yufeng Li 2023-05-11 13:47:15 -07:00 committed by GitHub
parent 27b2815d42
commit 902c5f53ae
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 302 additions and 67 deletions

View file

@ -858,6 +858,10 @@ Status QkvToContext(
query = data.query;
}
DUMP_TENSOR_D("attention q(BSNH)", q, batch_size * sequence_length, num_heads * qk_head_size);
DUMP_TENSOR_D("attention k(BSNH)", k, batch_size * sequence_length, num_heads * qk_head_size);
DUMP_TENSOR_D("attention v(BSNH)", v, batch_size * sequence_length, num_heads * v_head_size);
MemoryEfficientAttentionParams p;
p.sm = device_prop.major * 10 + device_prop.minor;
p.is_half = sizeof(T) == 2;
@ -881,7 +885,7 @@ Status QkvToContext(
p.workspace = MemoryEfficientAttentionParams::need_workspace(v_head_size, sizeof(T) == sizeof(float)) ? scratch1 : nullptr;
p.stream = stream;
run_memory_efficient_attention(p);
DUMP_TENSOR("cutlass output", data.output, batch_size * sequence_length, num_heads, v_head_size);
DUMP_TENSOR("attention cutlass output", data.output, batch_size * sequence_length, num_heads, v_head_size);
return Status::OK();
}
#endif

View file

@ -274,6 +274,20 @@ Status PackedAttention<T>::ComputeInternal(OpKernelContext* context) const {
MHARunner* fused_runner = TryGettingFusedRunner(parameters);
bool use_memory_efficient_attention = false;
auto& device_prop = GetDeviceProp();
#if USE_FLASH_ATTENTION
if (nullptr == fused_runner) {
int sm = device_prop.major * 10 + device_prop.minor;
bool is_good_for_rpb = !parameters.has_relative_position_bias || parameters.sequence_length % (4 * sizeof(T)) == 0;
use_memory_efficient_attention = is_good_for_rpb &&
sizeof(T) == 2 && // only enable for fp16
(parameters.head_size & 7) == 0 &&
(parameters.v_head_size & 7) == 0 &&
has_memory_efficient_attention(sm, sizeof(T) == 2);
}
#endif
typedef typename ToCudaType<T>::MappedType CudaT;
CudaT one = ToCudaType<T>::FromFloat(1.0f);
CudaT zero = ToCudaType<T>::FromFloat(0.0f);
@ -284,7 +298,6 @@ Status PackedAttention<T>::ComputeInternal(OpKernelContext* context) const {
int k = parameters.input_hidden_size;
gemm_buffer = GetScratchBuffer<T>(static_cast<size_t>(m) * n, context->GetComputeStream());
auto& device_prop = GetDeviceProp();
cublasHandle_t cublas = GetCublasHandle(context);
// Gemm, note that CUDA assumes col-major, so result(N, M) = 1 * weights x input + 1 x bias
@ -302,7 +315,8 @@ Status PackedAttention<T>::ComputeInternal(OpKernelContext* context) const {
parameters.head_size,
parameters.v_head_size,
parameters.sequence_length,
fused_runner);
fused_runner,
use_memory_efficient_attention);
auto work_space = GetScratchBuffer<void>(workSpaceSize, context->GetComputeStream());
typedef typename ToCudaType<T>::MappedType CudaT;
@ -315,6 +329,7 @@ Status PackedAttention<T>::ComputeInternal(OpKernelContext* context) const {
data.cumulative_sequence_length = cumulative_sequence_length->Data<int32_t>();
data.output = reinterpret_cast<CudaT*>(output->MutableData<T>());
data.fused_runner = reinterpret_cast<void*>(fused_runner);
data.use_memory_efficient_attention = use_memory_efficient_attention;
return QkvToContext<CudaT>(device_prop, cublas, Stream(context), parameters, data);
}

View file

@ -46,7 +46,8 @@ size_t GetAttentionWorkspaceSize(
size_t qk_head_size,
size_t v_head_size,
size_t sequence_length,
void* fused_runner) {
void* fused_runner,
bool use_memory_efficient_attention) {
// Note that q, k and v might need alignment for fused attention kernels.
const size_t qkv_bytes = element_size * batch_size * num_heads * sequence_length * (qk_head_size + qk_head_size + v_head_size);
@ -54,6 +55,19 @@ size_t GetAttentionWorkspaceSize(
return qkv_bytes;
}
#if USE_FLASH_ATTENTION
if (use_memory_efficient_attention) {
size_t fmha_buffer_bytes = 0;
if (MemoryEfficientAttentionParams::need_workspace(v_head_size, element_size == sizeof(float))) {
fmha_buffer_bytes = batch_size * sequence_length * num_heads * v_head_size * sizeof(float);
}
return qkv_bytes + fmha_buffer_bytes;
}
#else
ORT_UNUSED_PARAMETER(use_memory_efficient_attention);
#endif
return qkv_bytes + 2 * GetAttentionScratchSize(element_size, batch_size, num_heads, sequence_length);
}
@ -142,7 +156,7 @@ __global__ void AddBiasTransposeQKVPacked(
// Block: 256
// For memory efficient fMHA from CUTLASS. For future use, doesn't support fMHA from CUTLASS yet.
// Input: Tx3xNxH
// Output: 3xBxNxSxH
// Output: 3xTxNxH
// T is token_count
// B is batch_size
// S is sequence_length
@ -152,48 +166,27 @@ template <typename T>
__global__ void AddBiasTransposeQKVPackedCutlass(
const T* input,
const T* biases,
int32_t N,
int32_t H_QK,
int32_t H_V,
int32_t D_QK,
int32_t D_V,
T* q,
T* k,
T* v,
const int32_t* token_offset,
int32_t token_count) {
int s = blockIdx.x;
int b = blockIdx.y;
int token_idx = blockIdx.x;
int S = gridDim.x;
input += token_idx * (D_QK + D_QK + D_V);
q += token_idx * D_QK;
k += token_idx * D_QK;
v += token_idx * D_V;
const int packing_token_idx = b * S + s;
const int padding_token_idx = token_offset[packing_token_idx];
b = padding_token_idx / S;
s = padding_token_idx - b % S;
input += packing_token_idx * N * (H_QK + H_QK + H_V);
int k_offset = N * H_QK;
int v_offset = N * H_QK + N * H_QK;
q += (b * S * N + s * N) * H_QK;
k += (b * S * N + s * N) * H_QK;
v += (b * S * N + s * N) * H_V;
if (packing_token_idx < token_count) {
for (int i = threadIdx.x; i < N * H_QK; i += blockDim.x) {
if (token_idx < token_count) {
for (int i = threadIdx.x; i < D_QK; i += blockDim.x) {
q[i] = input[i] + biases[i];
k[i] = input[i + k_offset] + biases[i + k_offset];
k[i] = input[D_QK + i] + biases[D_QK + i];
}
for (int i = threadIdx.x; i < N * H_V; i += blockDim.x) {
v[i] = input[i + v_offset] + biases[i + v_offset];
}
} else {
for (int i = threadIdx.x; i < N * H_QK; i += blockDim.x) {
q[i] = biases[i];
k[i] = biases[i + k_offset];
}
for (int i = threadIdx.x; i < N * H_V; i += blockDim.x) {
v[i] = biases[i + v_offset];
for (int i = threadIdx.x; i < D_V; i += blockDim.x) {
v[i] = input[D_QK + D_QK + i] + biases[D_QK + D_QK + i];
}
}
}
@ -254,18 +247,16 @@ void InvokeAddBiasTranspose(
output + 2 * batch_size * sequence_length * num_heads * qk_head_size,
token_offset,
token_count);
} else if (format == AttentionQkvFormat::Q_K_V_BSNH) { // TODO: add memory efficient support
const dim3 grid(sequence_length, batch_size);
} else if (format == AttentionQkvFormat::Q_K_V_BSNH) {
const dim3 grid(token_count);
AddBiasTransposeQKVPackedCutlass<T><<<grid, kMAX_THREADS_PER_BLOCK, 0, stream>>>(
input,
biases,
num_heads,
qk_head_size,
v_head_size,
num_heads * qk_head_size,
num_heads * v_head_size,
output,
output + batch_size * sequence_length * num_heads * qk_head_size,
output + 2 * batch_size * sequence_length * num_heads * qk_head_size,
token_offset,
output + token_count * num_heads * qk_head_size,
output + 2 * token_count * num_heads * qk_head_size,
token_count);
} else {
ORT_ENFORCE(format == AttentionQkvFormat::QKV_BSN3H);
@ -381,15 +372,14 @@ Status LaunchTransposeRemovePadding(
const int batch_size, const int seq_len, const int number_heads, const int head_size,
cudaStream_t stream);
// input: [batch_size, number_heads, seq_len, head_size]
// output: [token_count, number_heads * head_size]
// input: [batch_size, number_heads, seq_len, head_size]
// output: [token_count, number_heads * head_size]
template <>
Status LaunchTransposeRemovePadding(
half* output, const half* input,
const int* token_offset, const int token_count,
const int batch_size, const int seq_len, const int number_heads, const int head_size,
cudaStream_t stream) {
// Make sure memory is aligned to 128 bit
ORT_ENFORCE(!(reinterpret_cast<size_t>(input) & 0xF) && !(reinterpret_cast<size_t>(output) & 0xF), "alignment");
@ -476,6 +466,74 @@ Status FusedScaledDotProductAttention(
return Status::OK();
}
#if USE_FLASH_ATTENTION
template <typename T>
Status FusedScaledDotProductAttentionCutlass(
const cudaDeviceProp& device_prop,
cudaStream_t stream,
PackedAttentionParameters& parameters,
PackedAttentionData<T>& data) {
const int batch_size = parameters.batch_size;
const int sequence_length = parameters.sequence_length;
const int num_heads = parameters.num_heads;
const int qk_head_size = parameters.head_size;
const int v_head_size = parameters.v_head_size;
LaunchAddBiasTranspose(data.gemm_buffer, data.bias, data.workspace,
batch_size, sequence_length,
num_heads, qk_head_size, v_head_size,
AttentionQkvFormat::Q_K_V_BSNH, data.token_offset,
parameters.token_count, stream);
DUMP_TENSOR_INIT();
DUMP_TENSOR_D("PackedAttention cutlass data.gemm_buffer", data.gemm_buffer, parameters.token_count, 3, num_heads * qk_head_size);
DUMP_TENSOR_D("PackedAttention cutlass data.bias", data.bias, 1, 3 * num_heads * qk_head_size);
// Q, K and V pointers
const int model_dimension_qk = num_heads * qk_head_size;
const int model_dimension_v = num_heads * v_head_size;
const size_t elements_qk = static_cast<size_t>(parameters.token_count) * static_cast<size_t>(model_dimension_qk);
const size_t elements_v = static_cast<size_t>(parameters.token_count) * static_cast<size_t>(model_dimension_v);
T* qkv = data.workspace;
T* query = qkv;
T* key = query + elements_qk;
T* value = key + elements_qk;
T* accum_workspace = value + elements_v;
DUMP_TENSOR_D("PackedAttention cutlass q(BSNH)", query, parameters.token_count, num_heads * qk_head_size);
DUMP_TENSOR_D("PackedAttention cutlass k(BSNH)", key, parameters.token_count, num_heads * qk_head_size);
DUMP_TENSOR_D("PackedAttention cutlass v(BSNH)", value, parameters.token_count, num_heads * v_head_size);
DUMP_TENSOR_D("PackedAttention cutlass cumulative_sequence_length", data.cumulative_sequence_length, 1, batch_size + 1);
MemoryEfficientAttentionParams p;
p.sm = device_prop.major * 10 + device_prop.minor;
p.is_half = sizeof(T) == 2;
p.batch_size = parameters.batch_size;
p.num_heads = parameters.num_heads;
p.sequence_length = parameters.sequence_length;
p.kv_sequence_length = parameters.sequence_length;
p.qk_head_size = parameters.head_size;
p.v_head_size = parameters.v_head_size;
p.causal = false;
p.scale = parameters.scale == 0.0f ? 1.f / sqrt(static_cast<float>(qk_head_size))
: parameters.scale;
p.seqlen_k_ptr = nullptr;
p.seqstart_q_ptr = const_cast<int32_t*>(data.cumulative_sequence_length);
p.seqstart_k_ptr = const_cast<int32_t*>(data.cumulative_sequence_length);
p.query = query;
p.key = key;
p.value = value;
p.attn_bias = data.relative_position_bias;
p.is_attn_bias_batched = !parameters.broadcast_res_pos_bias;
p.output = data.output;
p.workspace = MemoryEfficientAttentionParams::need_workspace(v_head_size, sizeof(T) == sizeof(float)) ? accum_workspace : nullptr;
p.stream = stream;
run_memory_efficient_attention(p);
DUMP_TENSOR("PackedAttention cutlass output", data.output, parameters.token_count, num_heads, v_head_size);
return Status::OK();
}
#endif
template <typename T>
Status UnfusedScaledDotProductAttention(
const cudaDeviceProp& device_prop,
@ -515,8 +573,8 @@ Status UnfusedScaledDotProductAttention(
// Q, K and V are ready now
DUMP_TENSOR_INIT();
DUMP_TENSOR_D("gemm_buffer", data.gemm_buffer, parameters.token_count, (num_heads * (qk_head_size * 2 + v_head_size)));
DUMP_TENSOR_D("data.workspace", data.workspace, 3 * batch_size, num_heads, sequence_length, qk_head_size);
DUMP_TENSOR_D("PackedAttention unfused gemm_buffer", data.gemm_buffer, parameters.token_count, (num_heads * (qk_head_size * 2 + v_head_size)));
DUMP_TENSOR_D("PackedAttention unfused data.workspace", data.workspace, 3 * batch_size, num_heads, sequence_length, qk_head_size);
// Compute Q*K' (as K'*Q), scaled by 1/sqrt(H) and store in scaled_qk: BxNxSxT
// Q: BxNxSxH, K: BxNxSxH, Q*K': BxNxSxS
@ -537,7 +595,7 @@ Status UnfusedScaledDotProductAttention(
scaled_qk, sequence_length, sequence_length * sequence_length,
batches, device_prop));
DUMP_TENSOR_D("QK", scaled_qk, batch_size * num_heads, sequence_length, sequence_length);
DUMP_TENSOR_D("PackedAttention unfused QK", scaled_qk, batch_size * num_heads, sequence_length, sequence_length);
const size_t bytes = GetAttentionScratchSize(element_size, batch_size, num_heads,
sequence_length);
@ -554,7 +612,7 @@ Status UnfusedScaledDotProductAttention(
num_heads,
attention_score, stream));
DUMP_TENSOR_D("Softmax", attention_score, batch_size * num_heads, sequence_length, sequence_length);
DUMP_TENSOR_D("PackedAttention unfused Softmax", attention_score, batch_size * num_heads, sequence_length, sequence_length);
// compute R*V (as V*R), and store in temp_output (space used by Q): BxNxSxH_v
T* temp_output = qkv;
@ -572,7 +630,7 @@ Status UnfusedScaledDotProductAttention(
batch_size, sequence_length, num_heads, v_head_size,
stream);
DUMP_TENSOR("unfused output", data.output, parameters.token_count, num_heads, v_head_size);
DUMP_TENSOR("PackedAttention unfused output", data.output, parameters.token_count, num_heads, v_head_size);
return result;
}
@ -586,9 +644,15 @@ Status QkvToContext(
void* fused_runner = data.fused_runner;
if (nullptr != fused_runner) {
return FusedScaledDotProductAttention<T>(device_prop, stream, parameters, data);
} else {
return UnfusedScaledDotProductAttention<T>(device_prop, cublas, stream, parameters, data);
}
#if USE_FLASH_ATTENTION
if (data.use_memory_efficient_attention) {
return FusedScaledDotProductAttentionCutlass(device_prop, stream, parameters, data);
}
#endif
return UnfusedScaledDotProductAttention<T>(device_prop, cublas, stream, parameters, data);
}
template Status QkvToContext<float>(

View file

@ -24,7 +24,8 @@ size_t GetAttentionWorkspaceSize(
size_t qk_head_size,
size_t v_head_size,
size_t sequence_length,
void* fused_runner);
void* fused_runner,
bool use_memory_efficient_attention);
template <typename T>
struct PackedAttentionData {
@ -38,6 +39,8 @@ struct PackedAttentionData {
T* output;
void* fused_runner;
bool use_memory_efficient_attention;
};
template <typename T>

View file

@ -385,14 +385,8 @@ static void RunModelWithRandomInput(
int64_t batch_size,
int64_t sequence_length,
std::string& onnx_model,
bool is_float16) {
// ORT enables TF32 in GEMM for A100. TF32 will cause precsion loss and fail this test.
// Do not run this test unless TF32 is disabled explicitly.
if (HasCudaEnvironment(800) && ParseEnvironmentVariableWithDefault<int>("NVIDIA_TF32_OVERRIDE", 1) != 0) {
GTEST_SKIP() << "Skipping RunModelWithRandomInput in A100 since TF32 is enabled";
return;
}
bool is_float16,
bool has_rbp = false) {
RandomValueGenerator random{234};
constexpr int hidden_size = 768;
@ -455,6 +449,16 @@ static void RunModelWithRandomInput(
test.AddInput<int32_t>("token_offset", token_offset_dims, token_offset);
test.AddInput<int32_t>("cumulative_sequence_length", cum_seq_len_dims, cum_seq_len);
if (has_rbp) {
std::vector<int64_t> rbp_dims{1, num_heads, sequence_length, sequence_length};
std::vector<float> rbp_data = random.Gaussian<float>(rbp_dims, 0.0f, 0.1f);
if (is_float16) {
test.AddInput<MLFloat16>("rbp", rbp_dims, ToFloat16(rbp_data));
} else {
test.AddInput<float>("rbp", rbp_dims, rbp_data);
}
}
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
execution_providers.push_back(DefaultCudaExecutionProvider());
test.AddReferenceOutputs(onnx_model, gpu_threshold, DefaultCudaExecutionProvider());
@ -462,7 +466,7 @@ static void RunModelWithRandomInput(
}
}
TEST(PackedAttentionTest, test_on_random_data) {
TEST(PackedAttentionTest, TestWithRandomData) {
std::string onnx_model = "testdata/packed_attention_fp32.onnx";
std::string onnx_model_fp16 = "testdata/packed_attention_fp16.onnx";
for (int batch_size : std::vector<int>({1, 2, 3, 4, 5, 6, 7, 8})) {
@ -481,7 +485,21 @@ TEST(PackedAttentionTest, test_on_random_data) {
}
}
TEST(PackedAttentionTest, test_on_random_data_large_seq) {
TEST(PackedAttentionTest, TestWithRandomDataWithRBP) {
std::string onnx_model_fp16 = "testdata/packed_attention_fp16.rbp.onnx"; // mainly for cutlass
for (int batch_size : std::vector<int>({1, 2, 3, 4, 5, 6, 7, 8})) {
for (int sequence_length : std::vector<int>({32, 48, 64, 95, 128})) {
RunModelWithRandomInput(
batch_size,
sequence_length,
onnx_model_fp16,
true /*is_float16*/,
true /*has_rbp*/);
}
}
}
TEST(PackedAttentionTest, TestWithRandomDataLargeSeq) {
int batch_size = 2;
int sequence_length = 1152; // > 1024
std::string onnx_model = "testdata/packed_attention_fp32.onnx";

Binary file not shown.

View file

@ -0,0 +1,131 @@
"""
Run this script to recreate the original onnx model.
Example usage:
python packed_attention_fp16.model.py out_model_path.onnx
"""
import sys
import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper
def clear_field(proto, field):
proto.ClearField(field)
return proto
def order_repeated_field(repeated_proto, key_name, order):
order = list(order)
repeated_proto.sort(key=lambda x: order.index(getattr(x, key_name)))
def make_node(op_type, inputs, outputs, name=None, doc_string=None, domain=None, **kwargs):
node = helper.make_node(op_type, inputs, outputs, name, doc_string, domain, **kwargs)
if not doc_string:
node.doc_string = ""
order_repeated_field(node.attribute, "name", kwargs.keys())
return node
def make_graph(*args, doc_string=None, **kwargs):
graph = helper.make_graph(*args, doc_string=doc_string, **kwargs)
if not doc_string:
graph.doc_string = ""
return graph
model = helper.make_model(
opset_imports=[
clear_field(helper.make_operatorsetid("", 12), "domain"),
helper.make_operatorsetid("com.microsoft", 1),
],
ir_version=7,
producer_name="onnxruntime.transformers",
producer_version="1.15.0",
graph=make_graph(
name="torch-jit-export",
inputs=[
helper.make_tensor_value_info("input", TensorProto.FLOAT16, shape=["token_count", "hidden_size"]),
helper.make_tensor_value_info("weight", TensorProto.FLOAT16, shape=["hidden_size", "hidden_size_x_3"]),
helper.make_tensor_value_info("bias", TensorProto.FLOAT16, shape=["hidden_size_x_3"]),
helper.make_tensor_value_info("token_offset", TensorProto.INT32, shape=["batch_size", "seq_len"]),
helper.make_tensor_value_info("cumulative_sequence_length", TensorProto.INT32, shape=["batch_size_plus_1"]),
helper.make_tensor_value_info(
"rbp", TensorProto.FLOAT16, shape=["batch_size", "num_heads", "seq_len", "seq_len"]
),
],
outputs=[helper.make_tensor_value_info("output", TensorProto.FLOAT16, shape=["token_count", "hidden_size"])],
nodes=[
make_node(
"Constant",
inputs=[],
outputs=["constant_0"],
name="constant_0",
value=numpy_helper.from_array(np.array([0], dtype="int64"), name=""),
),
make_node(
"Constant",
inputs=[],
outputs=["constant_1"],
name="Constant_1",
value=numpy_helper.from_array(np.array([1], dtype="int64"), name=""),
),
make_node(
"Constant",
inputs=[],
outputs=["constant_last"],
name="constant_last",
value=numpy_helper.from_array(np.array([-1], dtype="int64"), name=""),
),
make_node(
"Constant",
inputs=[],
outputs=["constant_max"],
name="Constant_max",
value=numpy_helper.from_array(np.array([9223372036854775807], dtype="int64"), name=""),
),
make_node(
"Slice",
inputs=["cumulative_sequence_length", "constant_0", "constant_last"],
outputs=["start"],
name="slice_start",
),
make_node(
"Slice",
inputs=["cumulative_sequence_length", "constant_1", "constant_max"],
outputs=["end"],
name="slice_end",
),
make_node("Sub", inputs=["end", "start"], outputs=["mask_idx"]),
make_node(
"RestorePadding",
inputs=["input", "token_offset"],
outputs=["restore_padding_output"],
name="RestorePadding_1",
domain="com.microsoft",
),
make_node(
"Attention",
inputs=["restore_padding_output", "weight", "bias", "mask_idx", "", "rbp"],
outputs=["attention_outputs"],
name="Attention_0",
domain="com.microsoft",
num_heads=12,
mask_filter_value=-3.4028234663852886e38,
),
make_node(
"RemovePadding",
inputs=["attention_outputs", "mask_idx"],
outputs=["output", "remove_padding_token_offset", "cumulated_seq_len", "max_seq_len"],
name="RemovePadding_0",
domain="com.microsoft",
),
],
),
)
if __name__ == "__main__" and len(sys.argv) == 2:
_, out_path = sys.argv
onnx.save(model, out_path)