From 436dfc3c9d44e069f8ad9645bfddaf51e6a2e47f Mon Sep 17 00:00:00 2001 From: Satya Kumar Jandhyala Date: Fri, 10 Jan 2025 13:31:50 -0800 Subject: [PATCH] [Native WebGPU] Fix the error when past and present key/value share buffer (#23315) ### Description Fix error causing incorrect output when past key/value share buffer with present key/value ### Motivation and Context --- .../contrib_ops/webgpu/bert/attention.cc | 29 ++++++++++--------- .../contrib_ops/webgpu/bert/attention.h | 14 +++++---- 2 files changed, 24 insertions(+), 19 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/attention.cc b/onnxruntime/contrib_ops/webgpu/bert/attention.cc index 86dc959cf2..568e75b38a 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/attention.cc @@ -69,10 +69,10 @@ Status TransferBSDToBNSH(onnxruntime::webgpu::ComputeContext& context, int num_h return context.RunProgram(program); }; -void InitVarStub(std::ostringstream& ss, const Tensor* seqlen_k, bool is_first_prompt) { +void InitVarStub(std::ostringstream& ss, const Tensor* seqlen_k) { if (seqlen_k != nullptr) { ss << "total_sequence_length = u32(seqlen_k[batch_idx]) + 1;\n"; - ss << "var past_sequence_length: u32 = " << (is_first_prompt ? "0" : "total_sequence_length - sequence_length") << ";\n"; + ss << "var past_sequence_length: u32 = select(total_sequence_length - sequence_length, 0u, uniforms.is_first_prompt > 0);\n"; } else { ss << "let past_sequence_length = uniforms.past_sequence_length;\n"; } @@ -106,7 +106,7 @@ Status AttentionProbsProgram::GenerateShaderCode(ShaderHelper& shader) const { << "let sequence_length = uniforms.M;\n" << "var total_sequence_length = uniforms.N;\n"; std::ostringstream oss; - InitVarStub(oss, seqlen_k_, is_first_prompt_); + InitVarStub(oss, seqlen_k_); shader.MainFunctionBody() << oss.str(); shader.MainFunctionBody() << "let kOffset = (workgroup_id.z / " << n_reps_ << ") * uniforms.kv_sequence_length * uniforms.K;\n"; if (has_present_key_) { @@ -121,7 +121,7 @@ Status AttentionProbsProgram::GenerateShaderCode(ShaderHelper& shader) const { " if (n + local_id.y < uniforms.N && w + local_id.x < uniforms.K) {\n" " var idx = TILE_SIZE * local_id.y + local_id.x;\n"; - if ((feed_past_key_ && has_present_key_) || past_present_share_buffer_) { + if ((feed_past_key_ && has_present_key_) || (past_present_share_buffer_ && !is_first_prompt_)) { shader.MainFunctionBody() << " if (n + local_id.y < past_sequence_length) {\n" << " let pastKeyOffset = (workgroup_id.z / " << n_reps_ << ") * uniforms.past_sequence_length * uniforms.K;\n" << " tileK[idx] = " << (past_present_share_buffer_ ? "present_key" : "past_key") << "[pastKeyOffset + (n + local_id.y) * uniforms.K + w + local_id.x];\n" @@ -213,7 +213,8 @@ Status ComputeAttentionProbs(onnxruntime::webgpu::ComputeContext& context, int o {static_cast(past_sequence_length)}, {static_cast(parameters.kv_sequence_length_)}, {static_cast(seqlen_k == nullptr ? total_sequence_length : parameters.seqlen_present_kv_cache_)}, - {static_cast(parameters.n_reps)}}) + {static_cast(parameters.n_reps)}, + {static_cast(parameters.is_first_prompt_ ? 1 : 0)}}) .SetOverridableConstants({{static_cast(tile_size)}}); return context.RunProgram(program); @@ -231,7 +232,7 @@ Status InPlaceSoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const { << "let sequence_length = uniforms.sequence_length;\n" << "var total_sequence_length = uniforms.total_sequence_length_comp * " << components_ << ";\n"; std::ostringstream oss; - InitVarStub(oss, seqlen_k_, is_first_prompt_); + InitVarStub(oss, seqlen_k_); shader.MainFunctionBody() << oss.str() << "let local_offset = local_idx * uniforms.elements_per_thread;\n" << "let offset = (global_idx / " << work_group_size_ << ") * uniforms.total_sequence_length_comp + local_offset;\n" @@ -285,12 +286,12 @@ Status ComputeInPlaceSoftmax(onnxruntime::webgpu::ComputeContext& context, Tenso } const int elementsPerThread = (total_sequence_length_comp + work_group_size - 1) / work_group_size; - InPlaceSoftmaxProgram program{"InPlaceSoftmax", work_group_size, components, is_first_prompt, seqlen_k}; + InPlaceSoftmaxProgram program{"InPlaceSoftmax", work_group_size, components, seqlen_k}; if (seqlen_k != nullptr) { program.AddInput({seqlen_k, ProgramTensorMetadataDependency::TypeAndRank}); } program.AddOutputs({{probs, ProgramTensorMetadataDependency::TypeAndRank, components}}) - .CacheHint(work_group_size, is_first_prompt) + .CacheHint(work_group_size) .SetDispatchGroupSize(1, sequence_length, batch_size * num_heads) .SetWorkgroupSize(work_group_size) .AddUniformVariables({{static_cast(batch_size)}, @@ -298,7 +299,8 @@ Status ComputeInPlaceSoftmax(onnxruntime::webgpu::ComputeContext& context, Tenso {static_cast(past_sequence_length)}, {static_cast(sequence_length)}, {static_cast(total_sequence_length_comp)}, - {static_cast(elementsPerThread)}}); + {static_cast(elementsPerThread)}, + {static_cast(is_first_prompt ? 1 : 0)}}); return context.RunProgram(program); } @@ -327,7 +329,7 @@ Status VxAttentionScoreProgram::GenerateShaderCode(ShaderHelper& shader) const { << "let sequence_length = uniforms.M;\n" << "var total_sequence_length = uniforms.K;\n"; std::ostringstream oss; - InitVarStub(oss, seqlen_k_, is_first_prompt_); + InitVarStub(oss, seqlen_k_); shader.MainFunctionBody() << oss.str(); shader.MainFunctionBody() << "let vOffset = (workgroup_id.z / " << n_reps_ << ") * uniforms.N * uniforms.kv_sequence_length + n;\n"; if (has_present_value_) { @@ -342,12 +344,12 @@ Status VxAttentionScoreProgram::GenerateShaderCode(ShaderHelper& shader) const { << " if (n < uniforms.N && w + local_id.y < uniforms.K) {\n" << " var idx = TILE_SIZE * local_id.y + local_id.x;\n"; - if ((feed_past_value_ && has_present_value_) || past_present_share_buffer_) { + if ((feed_past_value_ && has_present_value_) || (past_present_share_buffer_ && !is_first_prompt_)) { shader.MainFunctionBody() << " if (w + local_id.y < past_sequence_length) {\n" << " let pastValueOffset = (workgroup_id.z / " << n_reps_ << ") * uniforms.N * uniforms.past_sequence_length + n;\n" << " tileK[idx] = " << (past_present_share_buffer_ ? "present_value" : "past_value") << "[pastValueOffset + (w + local_id.y) * uniforms.N];\n" << " } else if (w + local_id.y - past_sequence_length < uniforms.kv_sequence_length) {\n" - << " tileK[idx] = v[vOffset + (w + local_id.y - uniforms.past_sequence_length) * uniforms.N];\n" + << " tileK[idx] = v[vOffset + (w + local_id.y - past_sequence_length) * uniforms.N];\n" << " }\n"; } else { shader.MainFunctionBody() << " if (w + local_id.y < uniforms.kv_sequence_length) {\n" @@ -425,7 +427,8 @@ Status ComputeVxAttentionScore(onnxruntime::webgpu::ComputeContext& context, int {static_cast(past_sequence_length)}, {static_cast(parameters.kv_sequence_length_)}, {static_cast(seqlen_k == nullptr ? total_sequence_length : parameters.seqlen_present_kv_cache_)}, - {static_cast(parameters.n_reps)}}) + {static_cast(parameters.n_reps)}, + {static_cast(parameters.is_first_prompt_)}}) .SetOverridableConstants({{static_cast(tile_size)}}); return context.RunProgram(program); diff --git a/onnxruntime/contrib_ops/webgpu/bert/attention.h b/onnxruntime/contrib_ops/webgpu/bert/attention.h index 03279fffbc..164ea72b07 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/attention.h @@ -49,7 +49,8 @@ class AttentionProbsProgram final : public Program { {"past_sequence_length", ProgramUniformVariableDataType::Uint32}, {"kv_sequence_length", ProgramUniformVariableDataType::Uint32}, {"present_sequence_length", ProgramUniformVariableDataType::Uint32}, - {"n_reps", ProgramUniformVariableDataType::Uint32}); + {"n_reps", ProgramUniformVariableDataType::Uint32}, + {"is_first_prompt", ProgramUniformVariableDataType::Uint32}); WEBGPU_PROGRAM_DEFINE_OVERRIDABLE_CONSTANTS({"TILE_SIZE", ProgramConstantDataType::Uint32}); @@ -67,8 +68,8 @@ class AttentionProbsProgram final : public Program { class InPlaceSoftmaxProgram final : public Program { public: - InPlaceSoftmaxProgram(const std::string& kernel_name, int work_group_size, int components, bool is_first_prompt, const Tensor* seqlen_k = nullptr) - : Program{kernel_name}, work_group_size_(work_group_size), components_(components), seqlen_k_(seqlen_k), is_first_prompt_(is_first_prompt) { + InPlaceSoftmaxProgram(const std::string& kernel_name, int work_group_size, int components, const Tensor* seqlen_k = nullptr) + : Program{kernel_name}, work_group_size_(work_group_size), components_(components), seqlen_k_(seqlen_k) { } Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -78,13 +79,13 @@ class InPlaceSoftmaxProgram final : public Program { {"past_sequence_length", ProgramUniformVariableDataType::Uint32}, {"sequence_length", ProgramUniformVariableDataType::Uint32}, {"total_sequence_length_comp", ProgramUniformVariableDataType::Uint32}, - {"elements_per_thread", ProgramUniformVariableDataType::Uint32}); + {"elements_per_thread", ProgramUniformVariableDataType::Uint32}, + {"is_first_prompt", ProgramUniformVariableDataType::Uint32}); private: int work_group_size_; int components_; const Tensor* seqlen_k_; - bool is_first_prompt_; }; class VxAttentionScoreProgram final : public Program { @@ -104,7 +105,8 @@ class VxAttentionScoreProgram final : public Program { {"past_sequence_length", ProgramUniformVariableDataType::Uint32}, {"kv_sequence_length", ProgramUniformVariableDataType::Uint32}, {"present_sequence_length", ProgramUniformVariableDataType::Uint32}, - {"n_reps", ProgramUniformVariableDataType::Uint32}); + {"n_reps", ProgramUniformVariableDataType::Uint32}, + {"is_first_prompt", ProgramUniformVariableDataType::Uint32}); WEBGPU_PROGRAM_DEFINE_OVERRIDABLE_CONSTANTS({"TILE_SIZE", ProgramConstantDataType::Uint32});