wehub-resource-sync
328 lines
12 KiB
Python
328 lines
12 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=invalid-name
|
|
# ruff: noqa: E501
|
|
"""Generator for CUTLASS attention kernels."""
|
|
|
|
from .library import substitute_template
|
|
|
|
|
|
def instantiate_attention_template(attrs):
|
|
"""Return CUTLASS host code for fused multi head attention
|
|
based on a template and the provided attribute map."""
|
|
|
|
bias_template = """
|
|
TVM_FFI_ICHECK(${bias}->ndim == 4); // B, N, S, S'
|
|
|
|
p.attn_bias_ptr = reinterpret_cast<T *>(${bias}->data);
|
|
p.bias_strideM = ${bias_strideM};
|
|
p.bias_strideH = ${bias_strideH};
|
|
p.bias_strideB = ${bias_strideB};
|
|
"""
|
|
|
|
var_len_template = """
|
|
p.seqstart_q_ptr = (int32_t*)${seqstart_q}->data;
|
|
p.seqstart_k_ptr = (int32_t*)${seqstart_k}->data;
|
|
p.num_queries = ((int32_t*)${max_seqlen_q}->data)[0];
|
|
p.num_batches = ${seqstart_q}->shape[0] - 1;
|
|
"""
|
|
|
|
qkv_template = {
|
|
"default": """
|
|
p.query_ptr = reinterpret_cast<T *>(${query}->data);
|
|
p.key_ptr = reinterpret_cast<T *>(${key}->data);
|
|
p.value_ptr = reinterpret_cast<T *>(${value}->data);
|
|
TVM_FFI_ICHECK(${query}->ndim == 4); // B, S, N, H
|
|
TVM_FFI_ICHECK(${key}->ndim == 4); // B, S', N, H
|
|
TVM_FFI_ICHECK(${value}->ndim == 4); // B, S', N, H'
|
|
|
|
// stride for N
|
|
p.q_strideH = p.head_dim; // H
|
|
p.k_strideH = p.head_dim; // H
|
|
p.v_strideH = p.head_dim_value; // H'
|
|
|
|
// stride for S
|
|
p.q_strideM = p.q_strideH * p.num_heads; // H * N
|
|
p.k_strideM = p.k_strideH * p.num_heads; // H * N
|
|
p.v_strideM = p.v_strideH * p.num_heads; // H' * N
|
|
|
|
// stride for B
|
|
p.q_strideB = p.q_strideM * p.num_queries; // H * N * S
|
|
p.k_strideB = p.k_strideM * p.num_keys; // H * N * S'
|
|
p.v_strideB = p.v_strideM * p.num_keys; // H'* N * S'
|
|
""",
|
|
"qkv_stacked": """
|
|
p.query_ptr = reinterpret_cast<T *>(${qkv}->data);
|
|
p.key_ptr = reinterpret_cast<T *>(${qkv}->data) + p.head_dim * p.num_heads;
|
|
p.value_ptr = reinterpret_cast<T *>(${qkv}->data) + p.head_dim * p.num_heads * 2;
|
|
TVM_FFI_ICHECK(${qkv}->ndim == 3); // B, S, NH + NH + NH'
|
|
|
|
// stride for N
|
|
p.q_strideH = p.head_dim; // H
|
|
p.k_strideH = p.head_dim; // H
|
|
p.v_strideH = p.head_dim_value; // H'
|
|
|
|
// stride for S
|
|
p.q_strideM = p.k_strideM = p.v_strideM =
|
|
p.q_strideH * p.num_heads +
|
|
p.k_strideH * p.num_heads +
|
|
p.v_strideH * p.num_heads; // H * N + H * N + H * N'
|
|
|
|
// stride for B
|
|
p.q_strideB = p.k_strideB = p.v_strideB =
|
|
p.q_strideM * p.num_queries; // (H * N + H * N + H * N') * S
|
|
""",
|
|
}
|
|
|
|
template = """
|
|
using T = ${data_type};
|
|
|
|
using Attention =
|
|
AttentionKernel<T,
|
|
/*ArchTag=*/${arch},
|
|
/*is_aligned=*/${kIsAligned},
|
|
/*queries_per_block=*/${kQueriesPerBlock},
|
|
/*keys_per_block=*/${kKeysPerBlock},
|
|
/*kMaxK=*/${kMaxK},
|
|
/*supports_dropout=*/${kSupportsDropout},
|
|
/*supports_bias=*/${kSupportsBias}
|
|
>;
|
|
|
|
typename Attention::Params p;
|
|
p.logsumexp_ptr = nullptr;
|
|
p.output_ptr = reinterpret_cast<T *>(out0->data);
|
|
|
|
p.output_accum_ptr = nullptr;
|
|
uint64_t accumulator_buf_size = ${output_size} * sizeof(Attention::output_accum_t);
|
|
bool accumulator_buf_allocated = false;
|
|
if (Attention::kNeedsOutputAccumulatorBuffer) {
|
|
if (accumulator_buf_size <= ${workspace}->shape[0]) {
|
|
p.output_accum_ptr = static_cast<float*>(${workspace}->data);
|
|
} else {
|
|
accumulator_buf_allocated = true;
|
|
cudaMalloc(
|
|
&p.output_accum_ptr,
|
|
accumulator_buf_size
|
|
);
|
|
}
|
|
}
|
|
|
|
p.num_heads = ${num_heads}; // N
|
|
p.num_batches = ${num_batches}; // B
|
|
p.head_dim = ${head_dim}; // H
|
|
p.head_dim_value = ${head_dim_value}; // H'
|
|
p.num_queries = ${num_queries}; // S
|
|
p.num_keys = ${num_keys}; // S'
|
|
p.scale = ${scale};
|
|
p.custom_mask_type = ${custom_mask_type};
|
|
|
|
|
|
p.o_strideM = p.head_dim_value * p.num_heads; // H' * N
|
|
TVM_FFI_ICHECK(out0->ndim == 4); // B, S, N, H'
|
|
|
|
${qkv_template}
|
|
${bias_template}
|
|
${var_len_template}
|
|
|
|
constexpr auto kernel_fn = attention_kernel_batched_impl<Attention>;
|
|
int smem_bytes = sizeof(typename Attention::SharedStorage);
|
|
if (smem_bytes > 0xc000) {
|
|
static bool once = [&]() {
|
|
cudaFuncSetAttribute(
|
|
kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
|
|
return true;
|
|
}();
|
|
}
|
|
|
|
TVM_FFI_ICHECK(Attention::check_supported(p));
|
|
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
|
|
|
|
kernel_fn<<<p.getBlocksGrid(), p.getThreadsGrid(), smem_bytes, stream>>>(p);
|
|
|
|
if (accumulator_buf_allocated) {
|
|
cudaFree(p.output_accum_ptr);
|
|
}
|
|
"""
|
|
|
|
template = substitute_template(
|
|
template,
|
|
{
|
|
"qkv_template": qkv_template[attrs["qkv_layout"]],
|
|
"bias_template": bias_template if "bias" in attrs else "",
|
|
"var_len_template": var_len_template if "seqstart_q" in attrs else "",
|
|
},
|
|
)
|
|
|
|
return substitute_template(template, attrs)
|
|
|
|
|
|
def instantiate_flash_attention_template(attrs):
|
|
"""Return host code for flash attention."""
|
|
|
|
template = """
|
|
int q_head_stride = ${head_dim};
|
|
int k_head_stride = ${head_dim};
|
|
int v_head_stride = ${head_dim};
|
|
int o_head_stride = ${head_dim};
|
|
int q_row_stride = q_head_stride * ${num_q_heads};
|
|
int k_row_stride = k_head_stride * ${num_kv_heads};
|
|
int v_row_stride = v_head_stride * ${num_kv_heads};
|
|
int o_row_stride = o_head_stride * ${num_q_heads};
|
|
int q_batch_stride = q_row_stride * ${num_queries};
|
|
int k_batch_stride = k_row_stride * ${num_keys};
|
|
int v_batch_stride = v_row_stride * ${num_keys};
|
|
int o_batch_stride = o_row_stride * ${num_queries};
|
|
|
|
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
|
|
|
|
flash_attn::flash_attention_forward(
|
|
static_cast<const cutlass::half_t*>(${query}->data),
|
|
static_cast<const cutlass::half_t*>(${key}->data),
|
|
static_cast<const cutlass::half_t*>(${value}->data),
|
|
static_cast<cutlass::half_t*>(out0->data),
|
|
${num_batches},
|
|
${num_queries},
|
|
${num_keys},
|
|
${num_q_heads},
|
|
${num_kv_heads},
|
|
${head_dim},
|
|
q_batch_stride,
|
|
k_batch_stride,
|
|
v_batch_stride,
|
|
o_batch_stride,
|
|
q_head_stride,
|
|
k_head_stride,
|
|
v_head_stride,
|
|
o_head_stride,
|
|
q_row_stride,
|
|
k_row_stride,
|
|
v_row_stride,
|
|
o_row_stride,
|
|
${scale},
|
|
${is_causal},
|
|
${window_size_left},
|
|
${window_size_right},
|
|
stream);
|
|
"""
|
|
|
|
template_stacked = """
|
|
int q_head_stride = ${head_dim};
|
|
int k_head_stride = ${head_dim};
|
|
int v_head_stride = ${head_dim};
|
|
int o_head_stride = ${head_dim};
|
|
int row_stride = q_head_stride * ${num_q_heads} +
|
|
k_head_stride * ${num_kv_heads} +
|
|
v_head_stride * ${num_kv_heads};
|
|
int q_row_stride = row_stride;
|
|
int k_row_stride = row_stride;
|
|
int v_row_stride = row_stride;
|
|
int o_row_stride = o_head_stride * ${num_q_heads};
|
|
|
|
int q_batch_stride = q_row_stride * ${num_queries};
|
|
int k_batch_stride = k_row_stride * ${num_keys};
|
|
int v_batch_stride = v_row_stride * ${num_keys};
|
|
int o_batch_stride = o_row_stride * ${num_queries};
|
|
|
|
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
|
|
|
|
flash_attn::flash_attention_forward(
|
|
static_cast<const cutlass::half_t*>(${qkv}->data),
|
|
static_cast<const cutlass::half_t*>(${qkv}->data) + ${head_dim} * ${num_q_heads},
|
|
static_cast<const cutlass::half_t*>(${qkv}->data) + ${head_dim} * (${num_q_heads} + ${num_kv_heads}),
|
|
static_cast<cutlass::half_t*>(out0->data),
|
|
${num_batches},
|
|
${num_queries},
|
|
${num_keys},
|
|
${num_q_heads},
|
|
${num_kv_heads},
|
|
${head_dim},
|
|
q_batch_stride,
|
|
k_batch_stride,
|
|
v_batch_stride,
|
|
o_batch_stride,
|
|
q_head_stride,
|
|
k_head_stride,
|
|
v_head_stride,
|
|
o_head_stride,
|
|
q_row_stride,
|
|
k_row_stride,
|
|
v_row_stride,
|
|
o_row_stride,
|
|
${scale},
|
|
${is_causal},
|
|
${window_size_left},
|
|
${window_size_right},
|
|
stream);
|
|
"""
|
|
|
|
if "qkv" in attrs:
|
|
return substitute_template(template_stacked, attrs)
|
|
|
|
return substitute_template(template, attrs)
|
|
|
|
|
|
def instantiate_flash_attention_var_len_template(attrs):
|
|
"""Return host code for flash attention with variable sequence lengths."""
|
|
|
|
template = """
|
|
int _max_seqlen_q = ((int32_t*)${max_seqlen_q}->data)[0];
|
|
int _max_seqlen_k = ((int32_t*)${max_seqlen_k}->data)[0];
|
|
|
|
int batch_size = ${seqstart_q}->shape[0] - 1;
|
|
|
|
int q_head_stride = ${head_dim};
|
|
int k_head_stride = ${head_dim};
|
|
int v_head_stride = ${head_dim};
|
|
int o_head_stride = ${head_dim};
|
|
int q_row_stride = q_head_stride * ${num_q_heads};
|
|
int k_row_stride = k_head_stride * ${num_kv_heads};
|
|
int v_row_stride = v_head_stride * ${num_kv_heads};
|
|
int o_row_stride = o_head_stride * ${num_q_heads};
|
|
|
|
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
|
|
|
|
flash_attn::flash_attention_var_len_forward(
|
|
static_cast<const cutlass::half_t*>(${query}->data),
|
|
static_cast<const cutlass::half_t*>(${key}->data),
|
|
static_cast<const cutlass::half_t*>(${value}->data),
|
|
static_cast<const int*>(${seqstart_q}->data),
|
|
static_cast<const int*>(${seqstart_k}->data),
|
|
static_cast<cutlass::half_t*>(out0->data),
|
|
batch_size,
|
|
_max_seqlen_q,
|
|
_max_seqlen_k,
|
|
${num_q_heads},
|
|
${num_kv_heads},
|
|
${head_dim},
|
|
q_head_stride,
|
|
k_head_stride,
|
|
v_head_stride,
|
|
o_head_stride,
|
|
q_row_stride,
|
|
k_row_stride,
|
|
v_row_stride,
|
|
o_row_stride,
|
|
${scale},
|
|
${is_causal},
|
|
// For SWA, is_causal must be false.
|
|
${is_causal} ? _max_seqlen_k : ${window_size_left},
|
|
${window_size_right},
|
|
stream);
|
|
"""
|
|
|
|
return substitute_template(template, attrs)
|