chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,327 @@
|
||||
# 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)
|
||||
Reference in New Issue
Block a user