saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Added split functionality to unpack packed-QKV.

Browse source
This commit is contained in:
Satya Kumar Jandhyala 2025-01-21 14:24:13 -08:00
parent 177f535bb2
commit 0b94f10276
2 changed files with 46 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc
View file

@ -30,6 +30,42 @@ ONNX_OPERATOR_KERNEL_EX(
.InputMemoryType(OrtMemTypeCPUInput, 6),
GroupQueryAttention);
Status SplitPackedQKVProgram::GenerateShaderCode(ShaderHelper& sh) const {
const auto& packed_qkv = sh.AddInput("packed_qkv", ShaderUsage::UseOffsetToIndices | ShaderUsage::UseUniform);
const auto& query = sh.AddOutput("query", ShaderUsage::UseSetByIndices | ShaderUsage::UseUniform);
const auto& key = sh.AddOutput("key", ShaderUsage::UseSetByIndices | ShaderUsage::UseUniform);
const auto& value = sh.AddOutput("val", ShaderUsage::UseSetByIndices | ShaderUsage::UseUniform);
sh.MainFunctionBody() << " let packed_qkv_indices = " << packed_qkv.OffsetToIndices("global_idx") << ";\n"
<< " let input_data = " << packed_qkv.GetByOffset("global_idx") << ";\n"
<< " let index = " << packed_qkv.IndicesGet("packed_qkv_indices", "2") << ";\n"
<< " if (index < uniforms.hidden_size) {\n"
<< " " << query.SetByIndices("packed_qkv_indices", "input_data") << ";\n"
<< " } else if (index < (uniforms.hidden_size + uniforms.kv_hidden_size)) {\n"
<< " var key_indices = packed_qkv_indices;\n"
<< " " << key.IndicesSet("key_indices", "2", "u32(index - uniforms.hidden_size)") << ";\n"
<< " " << key.SetByIndices("key_indices", "input_data") << ";\n"
<< " } else {\n"
<< " var val_indices = packed_qkv_indices;\n"
<< " " << value.IndicesSet("val_indices", "2", "u32(index - uniforms.hidden_size - uniforms.kv_hidden_size)") << ";\n"
<< " " << value.SetByIndices("val_indices", "input_data") << ";\n"
<< " }";
return Status::OK();
}
Status SplitPackedQKV(onnxruntime::webgpu::ComputeContext& context, const WebgpuAttentionParameters& params, const Tensor* packedQKV, Tensor* query, Tensor* key, Tensor* val) {
SplitPackedQKVProgram program(params);
auto input_size = packedQKV->Shape().Size();
program
.AddInput({packedQKV, ProgramTensorMetadataDependency::Rank})
.AddOutputs({{query, ProgramTensorMetadataDependency::Rank}, {key, ProgramTensorMetadataDependency::Rank}, {val, ProgramTensorMetadataDependency::Rank}})
.AddUniformVariables({
{static_cast<uint32_t>(params.hidden_size_)},
{static_cast<uint32_t>(params.kv_hidden_size_)},
})
.SetDispatchGroupSize((input_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE);
return context.RunProgram(program);
}
Status GeneratePositionIDsProgram::GenerateShaderCode(ShaderHelper& sh) const {
sh.AddInput("seqlens", ShaderUsage::UseUniform);
sh.AddOutput("output", ShaderUsage::UseUniform);

10
onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h
View file

@ -30,6 +30,16 @@ class GeneratePositionIDsProgram final : public Program<GeneratePositionIDsProgr
{"interleaved", ProgramUniformVariableDataType::Uint32},
{"is_first_prompt", ProgramUniformVariableDataType::Uint32},
{"total_seqlen", ProgramUniformVariableDataType::Uint32});
};
class SplitPackedQKVProgram final : public Program<SplitPackedQKVProgram> {
public:
SplitPackedQKVProgram(const WebgpuAttentionParameters& params) : Program{"SplitPackedQKV"}, params_(params) {}
Status GenerateShaderCode(ShaderHelper& sh) const override;
WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"hidden_size", ProgramUniformVariableDataType::Uint32},
{"kv_hidden_size", ProgramUniformVariableDataType::Uint32});
private:
const WebgpuAttentionParameters& params_;