mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Apply split and rotrary embedding before converting input ro BSD to BNSH
This commit is contained in:
Satya Kumar Jandhyala
parent
f0d238a3e8
commit
1009fc9a63
1 changed files with 60 additions and 68 deletions
|
|
@ -67,21 +67,17 @@ Status SplitPackedQKV(onnxruntime::webgpu::ComputeContext& context, const Webgpu
|
|||
}
|
||||
|
||||
Status GeneratePositionIDsProgram::GenerateShaderCode(ShaderHelper& sh) const {
|
||||
sh.AddInput("seqlens", ShaderUsage::UseUniform);
|
||||
sh.AddOutput("output", ShaderUsage::UseUniform);
|
||||
sh.MainFunctionBody() << "let batch_idx = global_idx / uniforms.sequence_length;\n"
|
||||
<< "let sequence_idx = i32(global_idx % uniforms.sequence_length);\n"
|
||||
<< "var pos_id: u32 = 0u;\n"
|
||||
<< "if (uniforms.is_first_prompt == 0) {\n"
|
||||
<< " let total_seqlen = seqlens[batch_idx] + 1;\n"
|
||||
const auto& output = sh.AddOutput("output", ShaderUsage::UseUniform);
|
||||
const auto& seqlens = sh.AddInput("seqlens", ShaderUsage::UseUniform);
|
||||
sh.MainFunctionBody() << " let batch_idx = global_idx / uniforms.sequence_length;\n"
|
||||
<< " let sequence_idx = i32(global_idx % uniforms.sequence_length);\n"
|
||||
<< " var pos_id: u32 = 1u;\n"
|
||||
<< " let total_seqlen = " << seqlens.GetByOffset("batch_idx") << " + 1;\n"
|
||||
<< " let past_seqlen = total_seqlen - i32(uniforms.sequence_length);\n"
|
||||
<< " if (past_seqlen + sequence_idx < total_seqlen) {\n"
|
||||
<< " pos_id = u32(past_seqlen + sequence_idx);\n"
|
||||
<< " } else {\n"
|
||||
<< " pos_id = 1u;\n"
|
||||
<< " }\n"
|
||||
<< "}\n"
|
||||
<< "output[global_idx] = vec2<u32>(pos_id, u32(select(0, 0xFFFFFFF, pos_id < 0)));\n";
|
||||
<< " }\n";
|
||||
sh.MainFunctionBody() << " " << output.SetByOffset("global_idx", "pos_id");
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
|
|
@ -96,10 +92,10 @@ Status GeneratePositionIDs(onnxruntime::webgpu::ComputeContext& context, const W
|
|||
return context.RunProgram(program);
|
||||
}
|
||||
|
||||
Status RunRotaryEmbedding(onnxruntime::webgpu::ComputeContext& context, const WebgpuAttentionParameters& params, const Tensor* input, const Tensor* pos_ids, const Tensor* cos_cache, const Tensor* sin_cache, Tensor* output, bool is_packed_qkv, bool is_query_input) {
|
||||
Status RunRotaryEmbedding(onnxruntime::webgpu::ComputeContext& context, const WebgpuAttentionParameters& params, const Tensor* input, const Tensor* pos_ids, const Tensor* cos_cache, const Tensor* sin_cache, Tensor* output, bool is_query_input) {
|
||||
const auto half_rotary_embedding_dim = gsl::narrow<uint32_t>(cos_cache->Shape()[1]);
|
||||
auto num_heads = is_packed_qkv ? params.num_heads_ + 2 * params.kv_num_heads_ : (is_query_input ? params.num_heads_ : params.kv_num_heads_);
|
||||
const TensorShape global_shape({params.batch_size_, params.sequence_length_, num_heads, static_cast<int64_t>(params.head_size_) - static_cast<int64_t>(half_rotary_embedding_dim)});
|
||||
const auto head_size = params.rotary_dim_ == 0 ? half_rotary_embedding_dim * 2 : params.head_size_;
|
||||
const TensorShape global_shape({params.batch_size_, params.sequence_length_, (is_query_input ? params.hidden_size_ : params.kv_hidden_size_) / head_size, static_cast<int64_t>(head_size) - static_cast<int64_t>(half_rotary_embedding_dim)});
|
||||
const auto rank = global_shape.NumDimensions();
|
||||
std::vector<uint32_t> global_dims(rank);
|
||||
std::vector<uint32_t> global_strides(rank);
|
||||
|
|
@ -107,7 +103,7 @@ Status RunRotaryEmbedding(onnxruntime::webgpu::ComputeContext& context, const We
|
|||
global_dims[j] = gsl::narrow<uint32_t>(global_shape[j]);
|
||||
global_strides[j] = gsl::narrow<uint32_t>(global_shape.SizeFromDimension(j + 1));
|
||||
}
|
||||
const auto input_output_strides = std::vector<uint32_t>({gsl::narrow<uint32_t>(params.batch_size_), gsl::narrow<uint32_t>(params.hidden_size_), gsl::narrow<uint32_t>(params.head_size_), 1});
|
||||
const auto input_output_strides = std::vector<uint32_t>({gsl::narrow<uint32_t>(input->Shape().SizeFromDimension(1)), gsl::narrow<uint32_t>(params.sequence_length_) * gsl::narrow<uint32_t>(head_size), head_size, 1});
|
||||
const auto output_size = gsl::narrow<const uint32_t>(global_shape.Size());
|
||||
|
||||
RotaryEmbeddingProgram program(params.rotary_interleaved_);
|
||||
|
|
@ -138,7 +134,7 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext&
|
|||
const Tensor* cos_cache = context.Input<Tensor>(7);
|
||||
const Tensor* sin_cache = context.Input<Tensor>(8);
|
||||
|
||||
GroupQueryAttentionParameters params;
|
||||
GroupQueryAttentionParameters params = {};
|
||||
ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckInputs(query,
|
||||
key,
|
||||
value,
|
||||
|
|
@ -170,60 +166,56 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext&
|
|||
Tensor* present_value = context.Output(2, present_kv_shape);
|
||||
parameters.past_present_share_buffer_ = present_key != nullptr && present_value != nullptr && past_key != nullptr && past_value != nullptr && past_key->DataRaw() == present_key->DataRaw() && past_value->DataRaw() == present_value->DataRaw();
|
||||
|
||||
TensorShapeVector q_new_dims({parameters.batch_size_, parameters.is_packed_qkv_ ? parameters.num_heads_ + 2 * parameters.kv_num_heads_ : parameters.num_heads_,
|
||||
parameters.sequence_length_, parameters.head_size_});
|
||||
Tensor qSplit;
|
||||
Tensor kSplit;
|
||||
Tensor vSplit;
|
||||
if (parameters.is_packed_qkv_) {
|
||||
qSplit = context.CreateGPUTensor(query->DataType(), TensorShape({parameters.batch_size_, parameters.sequence_length_, parameters.hidden_size_}));
|
||||
kSplit = context.CreateGPUTensor(query->DataType(), TensorShape({parameters.batch_size_, parameters.sequence_length_, parameters.kv_hidden_size_}));
|
||||
vSplit = context.CreateGPUTensor(query->DataType(), TensorShape({parameters.batch_size_, parameters.sequence_length_, parameters.kv_hidden_size_}));
|
||||
ORT_RETURN_IF_ERROR(SplitPackedQKV(context, parameters, query, &qSplit, &kSplit, &vSplit));
|
||||
parameters.is_packed_qkv_ = false;
|
||||
query = &qSplit;
|
||||
key = &kSplit;
|
||||
value = &vSplit;
|
||||
}
|
||||
|
||||
Tensor qRotary;
|
||||
Tensor kRotary;
|
||||
if (do_rotary_) {
|
||||
qRotary = context.CreateGPUTensor(query->DataType(), query->Shape());
|
||||
kRotary = context.CreateGPUTensor(key->DataType(), key->Shape());
|
||||
auto pos_ids_shape = parameters.is_first_prompt_ ? TensorShape({1, 1}) : TensorShape({parameters.batch_size_, parameters.sequence_length_});
|
||||
Tensor pos_ids = context.CreateGPUTensor(DataTypeImpl::GetType<int64_t>(), pos_ids_shape);
|
||||
if (!parameters.is_first_prompt_) {
|
||||
ORT_RETURN_IF_ERROR(GeneratePositionIDs(context, parameters, seqlens_k, &pos_ids));
|
||||
}
|
||||
ORT_RETURN_IF_ERROR(RunRotaryEmbedding(context, parameters, query, &pos_ids, cos_cache, sin_cache, &qRotary, /* is_query_input = */ true));
|
||||
|
||||
ORT_RETURN_IF_ERROR(RunRotaryEmbedding(context, parameters, key, &pos_ids, cos_cache, sin_cache, &kRotary, /* is_query_input = */ false));
|
||||
query = &qRotary;
|
||||
key = &kRotary;
|
||||
}
|
||||
|
||||
TensorShapeVector q_new_dims({parameters.batch_size_, parameters.num_heads_, parameters.sequence_length_, parameters.head_size_});
|
||||
TensorShape q_new_shape(q_new_dims);
|
||||
Tensor qBNSH = context.CreateGPUTensor(query->DataType(), q_new_shape);
|
||||
ORT_RETURN_IF_ERROR(TransferBSDToBNSH(
|
||||
context, parameters.num_heads_, parameters.sequence_length_, parameters.head_size_, query, nullptr, 0, &qBNSH));
|
||||
if (!parameters.is_packed_qkv_) {
|
||||
TensorShapeVector k_new_dims({parameters.batch_size_, parameters.kv_num_heads_,
|
||||
parameters.kv_sequence_length_, parameters.head_size_});
|
||||
TensorShape k_new_shape(k_new_dims);
|
||||
Tensor kBNSH = context.CreateGPUTensor(key->DataType(), k_new_shape);
|
||||
ORT_RETURN_IF_ERROR(TransferBSDToBNSH(context, parameters.kv_num_heads_, parameters.kv_sequence_length_,
|
||||
parameters.head_size_, key, nullptr, 0, &kBNSH));
|
||||
ORT_RETURN_IF_ERROR(TransferBSDToBNSH(context, parameters.num_heads_, parameters.sequence_length_, parameters.head_size_, query, nullptr, 0, &qBNSH));
|
||||
TensorShapeVector k_new_dims({parameters.batch_size_, parameters.kv_num_heads_,
|
||||
parameters.kv_sequence_length_, parameters.head_size_});
|
||||
TensorShape k_new_shape(k_new_dims);
|
||||
Tensor kBNSH = context.CreateGPUTensor(key->DataType(), k_new_shape);
|
||||
ORT_RETURN_IF_ERROR(TransferBSDToBNSH(context, parameters.kv_num_heads_, parameters.kv_sequence_length_,
|
||||
parameters.head_size_, key, nullptr, 0, &kBNSH));
|
||||
|
||||
TensorShapeVector v_new_dims({parameters.batch_size_, parameters.kv_num_heads_,
|
||||
parameters.kv_sequence_length_, parameters.v_head_size_});
|
||||
TensorShape v_new_shape(v_new_dims);
|
||||
Tensor vBNSH = context.CreateGPUTensor(value->DataType(), v_new_shape);
|
||||
ORT_RETURN_IF_ERROR(TransferBSDToBNSH(context, parameters.kv_num_heads_, parameters.kv_sequence_length_,
|
||||
parameters.v_head_size_, value, nullptr, 0, &vBNSH));
|
||||
|
||||
if (do_rotary_) {
|
||||
Tensor qRotary = context.CreateGPUTensor(qBNSH.DataType(), qBNSH.Shape());
|
||||
Tensor kRotary = context.CreateGPUTensor(kBNSH.DataType(), kBNSH.Shape());
|
||||
TensorShape pos_ids_shape = parameters.is_first_prompt_ ? TensorShape({1}) : TensorShape({static_cast<int64_t>(parameters.batch_size_) * static_cast<int64_t>(parameters.sequence_length_)});
|
||||
Tensor pos_ids = context.CreateGPUTensor(DataTypeImpl::GetType<int64_t>(), pos_ids_shape);
|
||||
ORT_RETURN_IF_ERROR(GeneratePositionIDs(context, parameters, seqlens_k, &pos_ids));
|
||||
|
||||
ORT_RETURN_IF_ERROR(RunRotaryEmbedding(context, parameters, &qBNSH, &pos_ids, cos_cache, sin_cache, &qRotary, /* is_packed_qkv = */ false, /* is_query_input = */ true));
|
||||
|
||||
ORT_RETURN_IF_ERROR(RunRotaryEmbedding(context, parameters, &kBNSH, &pos_ids, cos_cache, sin_cache, &kRotary, /* is_packed_qkv = */ false, /* is_query_input = */ false));
|
||||
return ApplyAttention(&qRotary, &kRotary, &vBNSH, nullptr, past_key, past_value, output, present_key,
|
||||
present_value, parameters, context, seqlens_k);
|
||||
} else {
|
||||
return ApplyAttention(&qBNSH, &kBNSH, &vBNSH, nullptr, past_key, past_value, output, present_key,
|
||||
present_value, parameters, context, seqlens_k);
|
||||
}
|
||||
} else {
|
||||
// Q, K and V are packed. Both key and value are nullptr
|
||||
if (do_rotary_) {
|
||||
Tensor qRotary = context.CreateGPUTensor(qBNSH.DataType(), qBNSH.Shape());
|
||||
TensorShape pos_ids_shape = parameters.is_first_prompt_ ? TensorShape({1}) : TensorShape({static_cast<int64_t>(parameters.batch_size_) * static_cast<int64_t>(parameters.sequence_length_)});
|
||||
Tensor pos_ids = context.CreateGPUTensor(DataTypeImpl::GetType<int64_t>(), pos_ids_shape);
|
||||
ORT_RETURN_IF_ERROR(GeneratePositionIDs(context, parameters, seqlens_k, &pos_ids));
|
||||
|
||||
ORT_RETURN_IF_ERROR(RunRotaryEmbedding(context, parameters, &qBNSH, &pos_ids, cos_cache, sin_cache, &qRotary, /* is_packed_qkv = */ true, /* is_query_input = */ true));
|
||||
|
||||
return ApplyAttention(&qRotary, nullptr, nullptr, nullptr, past_key, past_value, output, present_key,
|
||||
present_value, parameters, context, seqlens_k);
|
||||
} else {
|
||||
return ApplyAttention(&qBNSH, nullptr, nullptr, nullptr, past_key, past_value, output, present_key,
|
||||
present_value, parameters, context, seqlens_k);
|
||||
}
|
||||
}
|
||||
TensorShapeVector v_new_dims({parameters.batch_size_, parameters.kv_num_heads_,
|
||||
parameters.kv_sequence_length_, parameters.v_head_size_});
|
||||
TensorShape v_new_shape(v_new_dims);
|
||||
Tensor vBNSH = context.CreateGPUTensor(value->DataType(), v_new_shape);
|
||||
ORT_RETURN_IF_ERROR(TransferBSDToBNSH(context, parameters.kv_num_heads_, parameters.kv_sequence_length_,
|
||||
parameters.v_head_size_, value, nullptr, 0, &vBNSH));
|
||||
return ApplyAttention(&qBNSH, &kBNSH, &vBNSH, nullptr, past_key, past_value, output, present_key,
|
||||
present_value, parameters, context, seqlens_k);
|
||||
}
|
||||
|
||||
} // namespace webgpu
|
||||
|
|
|
|||
Loading…
Reference in a new issue