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

Removed unnecessary uniforms in GeneratePositionIdsProgram

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

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

@ -86,18 +86,12 @@ Status GeneratePositionIDsProgram::GenerateShaderCode(ShaderHelper& sh) const {
}
Status GeneratePositionIDs(onnxruntime::webgpu::ComputeContext& context, const WebgpuAttentionParameters& params, const Tensor* seqlens, Tensor* output_tensor) {
GeneratePositionIDsProgram program(params);
auto output_size = params.is_first_prompt_ ? 1 : params.batch_size_ * params.sequence_length_;
program.AddInput(seqlens)
.AddOutput(output_tensor)
.AddUniformVariables({{static_cast<uint32_t>(params.batch_size_)},
{static_cast<uint32_t>(params.sequence_length_)},
{static_cast<uint32_t>(params.num_heads_)},
{static_cast<uint32_t>(params.head_size_)},
{static_cast<uint32_t>(params.rotary_dim_)},
{static_cast<uint32_t>(params.rotary_interleaved_)},
{static_cast<uint32_t>(params.is_first_prompt_ ? 0 : 1)},
{static_cast<uint32_t>(params.total_sequence_length_)}})
GeneratePositionIDsProgram program;
auto output_size = params.batch_size_ * params.sequence_length_;
program.CacheHint(params.batch_size_, params.sequence_length_)
.AddInput({seqlens, ProgramTensorMetadataDependency::Rank})
.AddOutput({output_tensor, ProgramTensorMetadataDependency::Rank})
.AddUniformVariables({{static_cast<uint32_t>(params.sequence_length_)}})
.SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE);
return context.RunProgram(program);
}

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

@ -18,18 +18,11 @@ Status GeneratePositionIDs(onnxruntime::webgpu::ComputeContext& context, bool is
class GeneratePositionIDsProgram final : public Program<GeneratePositionIDsProgram> {
public:
GeneratePositionIDsProgram(const WebgpuAttentionParameters& params) : Program{"GeneratePositionIDs"}, params_(params) {}
GeneratePositionIDsProgram() : Program{"GeneratePositionIDs"} {}
Status GenerateShaderCode(ShaderHelper& sh) const override;
WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"batch_size", ProgramUniformVariableDataType::Uint32},
{"sequence_length", ProgramUniformVariableDataType::Uint32},
{"num_heads", ProgramUniformVariableDataType::Uint32},
{"head_size", ProgramUniformVariableDataType::Uint32},
{"rotary_embedding_dim", ProgramUniformVariableDataType::Uint32},
{"interleaved", ProgramUniformVariableDataType::Uint32},
{"is_first_prompt", ProgramUniformVariableDataType::Uint32},
{"total_seqlen", ProgramUniformVariableDataType::Uint32});
WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"sequence_length", ProgramUniformVariableDataType::Uint32});
};
class SplitPackedQKVProgram final : public Program<SplitPackedQKVProgram> {