diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index b2979a95ce22e..e4e43535e2cfc 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -653,7 +653,8 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co Q, seqlen_k, cos_cache, sin_cache, &query_output, tq_present_key, tq_present_value, - indirect_buffer_ptr, tile_size, num_q_tiles)); + indirect_buffer_ptr, tile_size, num_q_tiles, + total_seqlen)); } else { ORT_RETURN_IF_ERROR(RunSplitPackedQKVWithRotaryEmbeddingAndCopyKV(context, parameters, Q, seqlen_k, @@ -668,7 +669,8 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co ORT_ENFORCE(K != nullptr && V != nullptr, "TurboQuant requires non-null K/V inputs when kv_sequence_length > 0."); ORT_RETURN_IF_ERROR(TurboQuantCopyToQuantizedKVCache(context, parameters, K, tq_past_key, tq_present_key, V, tq_past_value, tq_present_value, - tile_size, use_seqlen_k ? seqlen_k : nullptr, indirect_buffer_ptr, num_q_tiles)); + tile_size, use_seqlen_k ? seqlen_k : nullptr, indirect_buffer_ptr, num_q_tiles, + total_seqlen)); } else { ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, tile_size, use_seqlen_k ? seqlen_k : nullptr, indirect_buffer_ptr, num_q_tiles, total_seqlen)); } @@ -893,6 +895,7 @@ Status RunSplitPackedQKVWithRotaryEmbeddingAndCopyKV(onnxruntime::webgpu::Comput {static_cast(dispatch_size)}, {static_cast(params.batch_size_)}, {num_q_tiles}, + {static_cast(params.total_sequence_length_)}, }); program.SetDispatchGroupSize((dispatch_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE); diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index f917f50c7fd72..8ec7b57d445fe 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -37,7 +37,8 @@ class SplitPackedQKVWithRotaryEmbeddingAndCopyKVProgram final : public Program 0), + WGSL_TEMPLATE_PARAMETER(use_total_sequence_length_input, use_total_sequence_length_input_), WGSL_TEMPLATE_VARIABLE(cos_cache, cos_cache), WGSL_TEMPLATE_VARIABLE(key, key), WGSL_TEMPLATE_VARIABLE(packed_qkv, packed_qkv), @@ -50,6 +54,7 @@ Status RunSplitPackedQKVWithRotaryEmbedding(onnxruntime::webgpu::ComputeContext& const WebgpuAttentionParameters& params, const Tensor* packedQKV, const Tensor* seqlen_k, + const Tensor* total_seqlen, const Tensor* cos_cache, const Tensor* sin_cache, Tensor* query, @@ -79,15 +84,23 @@ Status RunSplitPackedQKVWithRotaryEmbedding(onnxruntime::webgpu::ComputeContext& auto dispatch_size = static_cast(params.batch_size_ * params.sequence_length_ * params.num_heads_ * work_per_head_vec); const uint32_t multi_rotary_cache_concat_offset = context.MultiRotaryCacheConcatOffset(); - SplitPackedQKVWithRotaryEmbeddingProgram program(params.rotary_interleaved_, multi_rotary_cache_concat_offset); + const bool use_total_sequence_length_input = + context.IsGraphCaptureEnabled() && multi_rotary_cache_concat_offset > 0; + SplitPackedQKVWithRotaryEmbeddingProgram program(params.rotary_interleaved_, + multi_rotary_cache_concat_offset, + use_total_sequence_length_input); program - .CacheHint(params.rotary_interleaved_, multi_rotary_cache_concat_offset) + .CacheHint(params.rotary_interleaved_, multi_rotary_cache_concat_offset, use_total_sequence_length_input) .AddInput({packedQKV, ProgramTensorMetadataDependency::TypeAndRank, components}) .AddInputs({ {seqlen_k, ProgramTensorMetadataDependency::TypeAndRank}, {cos_cache, ProgramTensorMetadataDependency::Rank, components}, {sin_cache, ProgramTensorMetadataDependency::Rank, components}, - }) + }); + if (use_total_sequence_length_input) { + program.AddInput({total_seqlen, ProgramTensorMetadataDependency::None}); + } + program .AddOutputs({{query, ProgramTensorMetadataDependency::None, components}, {key, ProgramTensorMetadataDependency::None, components}, {val, ProgramTensorMetadataDependency::None, components}}) @@ -99,6 +112,7 @@ Status RunSplitPackedQKVWithRotaryEmbedding(onnxruntime::webgpu::ComputeContext& {static_cast(params.kv_num_heads_)}, {static_cast(head_size_vec)}, {static_cast(half_rotary_embedding_dim_vec)}, + {static_cast(params.total_sequence_length_)}, {static_cast(dispatch_size)}, }) .SetDispatchGroupSize((dispatch_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE); @@ -419,7 +433,7 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& 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(RunSplitPackedQKVWithRotaryEmbedding(context, parameters, - query, seqlen_k, + query, seqlen_k, total_seqlen_tensor, cos_cache, sin_cache, &qSplit, &kSplit, &vSplit)); parameters.is_packed_qkv_ = false; diff --git a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h index 8af15937001a2..e232c06b3b4de 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h @@ -16,10 +16,13 @@ using namespace onnxruntime::webgpu; class SplitPackedQKVWithRotaryEmbeddingProgram final : public Program { public: - SplitPackedQKVWithRotaryEmbeddingProgram(bool interleaved, uint32_t multi_rotary_cache_concat_offset) + SplitPackedQKVWithRotaryEmbeddingProgram(bool interleaved, + uint32_t multi_rotary_cache_concat_offset, + bool use_total_sequence_length_input) : Program{"SplitPackedQKVWithRotaryEmbedding"}, interleaved_{interleaved}, - multi_rotary_cache_concat_offset_{multi_rotary_cache_concat_offset} {} + multi_rotary_cache_concat_offset_{multi_rotary_cache_concat_offset}, + use_total_sequence_length_input_{use_total_sequence_length_input} {} Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -31,11 +34,13 @@ class SplitPackedQKVWithRotaryEmbeddingProgram final : public Program multi_rotary_cache_concat_offset); + let base_position = select(0u, multi_rotary_cache_concat_offset, global_total_seq_length > multi_rotary_cache_concat_offset); #else let base_position = 0u; #endif diff --git a/onnxruntime/contrib_ops/webgpu/bert/split_packed_qkv_with_rotary_embedding_and_copykv.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/split_packed_qkv_with_rotary_embedding_and_copykv.wgsl.template index 6d88f883a5abb..487ca7c3416c8 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/split_packed_qkv_with_rotary_embedding_and_copykv.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/bert/split_packed_qkv_with_rotary_embedding_and_copykv.wgsl.template @@ -33,21 +33,25 @@ $MAIN { // Calculate position_id (needed for rotary embedding) let seqlen_i = seqlens.getByOffset(batch_idx); let seqlen = u32(seqlen_i); - let total_seqlen = seqlen + 1u; + let per_batch_total_seq_length = seqlen + 1u; // Right-padded batches with prompt shorter than sequence_length would underflow u32; clamp to 0. - let past_seqlen = select(total_seqlen - uniforms.sequence_length, 0u, total_seqlen <= uniforms.sequence_length); + let past_seqlen = per_batch_total_seq_length - min(per_batch_total_seq_length, uniforms.sequence_length); // `position_id` is used to get cos/sin cache and also as the time step index in present_key/present_value let position_id = past_seqlen + seq_idx; +#if prepare_indirect_dispatch + let global_total_seq_length = u32(total_sequence_length_input[0]); +#else + let global_total_seq_length = uniforms.total_sequence_length; +#endif #if use_multi_rotary_cache_concat - let base_position = select(0u, multi_rotary_cache_concat_offset, total_seqlen > multi_rotary_cache_concat_offset); + let base_position = select(0u, multi_rotary_cache_concat_offset, global_total_seq_length > multi_rotary_cache_concat_offset); #else let base_position = 0u; #endif #if prepare_indirect_dispatch if (global_idx == 0u) { - let global_total_seq_length = u32(total_sequence_length_input[0]); let num_total_seq_length_tile = (global_total_seq_length + uniforms.tile_size - 1u) / uniforms.tile_size; populate_indirect_dispatch_buffer(num_total_seq_length_tile, uniforms.num_heads * uniforms.num_q_tiles, uniforms.batch_size); } diff --git a/onnxruntime/contrib_ops/webgpu/bert/turbo_quant_fused_rotary_hadamard.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/turbo_quant_fused_rotary_hadamard.wgsl.template index 5c452b5ea188b..0f87349c4f661 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/turbo_quant_fused_rotary_hadamard.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/bert/turbo_quant_fused_rotary_hadamard.wgsl.template @@ -42,18 +42,39 @@ var scale_reduction_buffer : array; var index_buffer : array; $MAIN { - // Compute total_seq_length. + // Map flat workgroup index to logical component (Q, K, or V). + let num_kv_slices = uniforms.num_kv_slices; + let is_q = workgroup_idx >= 2u * num_kv_slices; + let is_value = !is_q && workgroup_idx >= num_kv_slices; + + // Compute batch based on workgroup type. + var batch: u32; + if (is_q) { + let q_slice = workgroup_idx - 2u * num_kv_slices; + if (q_slice >= uniforms.num_q_slices) { return; } + batch = q_slice / (uniforms.kv_sequence_length * uniforms.num_heads); + } else { + let kv_slice = select(workgroup_idx, workgroup_idx - num_kv_slices, is_value); + if (kv_slice >= num_kv_slices) { return; } + batch = kv_slice / (uniforms.kv_num_heads * uniforms.kv_sequence_length); + } + + // Compute the logical total sequence length for this batch. #if use_seqlen_k - let total_seq_length = u32(seqlen_k[0u]) + 1u; + let per_batch_total_seq_length = u32(seqlen_k[batch]) + 1u; #else - let total_seq_length = uniforms.total_sequence_length; + let per_batch_total_seq_length = uniforms.total_sequence_length; #endif - let past_seq_length = total_seq_length - uniforms.kv_sequence_length; + let past_seq_length = per_batch_total_seq_length - min(per_batch_total_seq_length, uniforms.kv_sequence_length); // Base position offset for rotary embedding cos/sin cache lookup. - let position_id = past_seq_length; +#if prepare_indirect_dispatch + let global_total_seq_length = u32(total_sequence_length_input[0]); +#else + let global_total_seq_length = uniforms.total_sequence_length; +#endif #if use_multi_rotary_cache_concat - let base_position = select(0u, multi_rotary_cache_concat_offset, total_seq_length > multi_rotary_cache_concat_offset); + let base_position = select(0u, multi_rotary_cache_concat_offset, global_total_seq_length > multi_rotary_cache_concat_offset); #else let base_position = 0u; #endif @@ -61,23 +82,18 @@ $MAIN { // Prepare indirect dispatch buffer (first workgroup, first thread only). #if prepare_indirect_dispatch if (workgroup_idx == 0u && local_idx == 0u) { - let num_total_seq_length_tile = (total_seq_length + uniforms.tile_size - 1u) / uniforms.tile_size; + let num_total_seq_length_tile = + (global_total_seq_length + uniforms.tile_size - 1u) / uniforms.tile_size; populate_indirect_dispatch_buffer(num_total_seq_length_tile, uniforms.num_heads * uniforms.num_q_tiles, uniforms.batch_size); } #endif - // Map flat workgroup index to logical component (Q, K, or V). - let num_kv_slices = uniforms.num_kv_slices; - let is_q = workgroup_idx >= 2u * num_kv_slices; - let is_value = !is_q && workgroup_idx >= num_kv_slices; - // ============ Q WORKGROUP PATH (early return, no shared memory/barriers) ============ if (is_q) { let q_slice = workgroup_idx - 2u * num_kv_slices; if (q_slice >= uniforms.num_q_slices) { return; } // Unflatten q_slice into (batch, seq, head). - let batch = q_slice / (uniforms.kv_sequence_length * uniforms.num_heads); let head = (q_slice / uniforms.kv_sequence_length) % uniforms.num_heads; let seq = q_slice % uniforms.kv_sequence_length; @@ -120,11 +136,15 @@ $MAIN { let kv_slice = select(workgroup_idx, workgroup_idx - num_kv_slices, is_value); if (kv_slice >= num_kv_slices) { return; } - // Unflatten kv_slice into (batch, head, seq). - let batch = kv_slice / (uniforms.kv_num_heads * uniforms.kv_sequence_length); + // Unflatten kv_slice into (head, seq) — batch already computed above. let head = (kv_slice / uniforms.kv_sequence_length) % uniforms.kv_num_heads; let seq = kv_slice % uniforms.kv_sequence_length; + // Skip K/V slices beyond this batch's logical total sequence length. + if (seq >= per_batch_total_seq_length) { + return; + } + // Compute destination offset in present_key/present_value (u32 packed, BNSH layout). #if past_present_share_buffer let dest_seq = past_seq_length + seq; diff --git a/onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.cc b/onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.cc index ab251d1a033fb..d0cbfe492ef34 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.cc @@ -24,7 +24,9 @@ Status TurboQuantHadamardProgram::GenerateShaderCode(ShaderHelper& shader) const if (use_seqlen_k_) { shader.AddInput("seqlen_k", ShaderUsage::None); } + // Use the batch-wide GPU value for dispatch sizing: batch 0 is not necessarily the longest batch. if (prepare_indirect_dispatch_) { + shader.AddInput("total_sequence_length_input", ShaderUsage::None); shader.AddOutput("indirect_buffer", ShaderUsage::None); } @@ -60,7 +62,7 @@ Status TurboQuantCopyToQuantizedKVCache(onnxruntime::webgpu::ComputeContext& con const Tensor* K, const Tensor* past_key, Tensor* present_key, const Tensor* V, const Tensor* past_value, Tensor* present_value, uint32_t tile_size, const Tensor* seqlen_k, Tensor* indirect_buffer, - uint32_t num_q_tiles) { + uint32_t num_q_tiles, const Tensor* total_seqlen) { const int head_size = parameters.head_size_; const int components = head_size % 4 == 0 ? 4 : (head_size % 2 == 0 ? 2 : 1); ORT_ENFORCE((head_size & (head_size - 1)) == 0 && head_size >= 8, @@ -81,13 +83,6 @@ Status TurboQuantCopyToQuantizedKVCache(onnxruntime::webgpu::ComputeContext& con bool prepare_indirect_dispatch = (indirect_buffer != nullptr); bool use_seqlen_k = (seqlen_k != nullptr); - ORT_RETURN_IF_ERROR( - (!use_seqlen_k || parameters.batch_size_ == 1) - ? Status::OK() - : ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "TurboQuant graph-capture decode path reads seqlen_k[0] for all batches and " - "currently supports batch_size == 1 only; got batch_size = ", - parameters.batch_size_)); bool kv_BNSH = parameters.qkv_format_ == Q_K_V_BSNH_BNSH_BNSH || parameters.qkv_format_ == Q_K_V_BNSH; TurboQuantHadamardProgram program{"TurboQuantCopyToQuantizedKVCache", has_past, kv_BNSH, @@ -113,6 +108,9 @@ Status TurboQuantCopyToQuantizedKVCache(onnxruntime::webgpu::ComputeContext& con if (use_seqlen_k) { program.AddInput({seqlen_k, ProgramTensorMetadataDependency::None}); } + if (prepare_indirect_dispatch) { + program.AddInput({total_seqlen, ProgramTensorMetadataDependency::None}); + } // Past KV cache is already u32-packed (no vectorization). if (has_past) { @@ -128,6 +126,7 @@ Status TurboQuantCopyToQuantizedKVCache(onnxruntime::webgpu::ComputeContext& con program.AddOutput({indirect_buffer, ProgramTensorMetadataDependency::None}); } + const uint32_t past_input_seq_length = has_past ? static_cast(past_key->Shape()[2]) : 0u; // present_key has shape (batch, kv_num_heads, present_seq_length, compressed_head_size_u32) uint32_t present_seq_length = static_cast(present_key->Shape()[2]); @@ -137,11 +136,13 @@ Status TurboQuantCopyToQuantizedKVCache(onnxruntime::webgpu::ComputeContext& con prepare_indirect_dispatch, use_seqlen_k, head_size_log2, components, compressed_head_size_u32) .AddUniformVariables({{static_cast(parameters.batch_size_)}, {static_cast(compressed_head_size_u32)}, + {static_cast(copy_sequence_length)}, {static_cast(kv_num_heads)}, {static_cast(parameters.kv_sequence_length_)}, {static_cast(parameters.num_heads_)}, {num_q_tiles}, {num_slices_per_kv}, + {past_input_seq_length}, {present_seq_length}, {tile_size}, {static_cast(parameters.total_sequence_length_)}}); @@ -157,6 +158,10 @@ Status TurboQuantFusedRotaryProgram::GenerateShaderCode(ShaderHelper& shader) co if (use_seqlen_k_) { shader.AddInput("seqlen_k", ShaderUsage::None); } + // Use the batch-wide GPU value for dispatch sizing: batch 0 is not necessarily the longest batch. + if (prepare_indirect_dispatch_) { + shader.AddInput("total_sequence_length_input", ShaderUsage::None); + } const auto& query = shader.AddOutput("query", ShaderUsage::UseUniform); // present_key/present_value are u32 arrays (packed 4-bit quantized data) @@ -195,7 +200,8 @@ Status TurboQuantApplyRotaryAndCopyToQuantizedKVCache(onnxruntime::webgpu::Compu Tensor* present_value, Tensor* indirect_buffer, uint32_t tile_size, - uint32_t num_q_tiles) { + uint32_t num_q_tiles, + const Tensor* total_seqlen) { const int head_size = parameters.head_size_; ORT_ENFORCE((head_size & (head_size - 1)) == 0 && head_size >= 8, "head_size must be a power of 2 >= 8 for TurboQuant fused rotary, got ", head_size); @@ -215,13 +221,6 @@ Status TurboQuantApplyRotaryAndCopyToQuantizedKVCache(onnxruntime::webgpu::Compu bool prepare_indirect_dispatch = (indirect_buffer != nullptr); bool use_seqlen_k = (seqlen_k != nullptr); - ORT_RETURN_IF_ERROR( - (!use_seqlen_k || parameters.batch_size_ == 1) - ? Status::OK() - : ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "TurboQuant graph-capture decode path reads seqlen_k[0] for all batches and " - "currently supports batch_size == 1 only; got batch_size = ", - parameters.batch_size_)); const uint32_t multi_rotary_cache_concat_offset = context.MultiRotaryCacheConcatOffset(); TurboQuantFusedRotaryProgram program{"TurboQuantFusedRotary", head_size_log2, @@ -240,6 +239,9 @@ Status TurboQuantApplyRotaryAndCopyToQuantizedKVCache(onnxruntime::webgpu::Compu if (use_seqlen_k) { program.AddInput({seqlen_k, ProgramTensorMetadataDependency::None}); } + if (prepare_indirect_dispatch) { + program.AddInput({total_seqlen, ProgramTensorMetadataDependency::None}); + } program.AddOutputs({{query, ProgramTensorMetadataDependency::None}, {present_key, ProgramTensorMetadataDependency::Rank}, diff --git a/onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.h b/onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.h index 4bea6f8237094..874ea23c06731 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.h +++ b/onnxruntime/contrib_ops/webgpu/bert/turbo_quant_hadamard.h @@ -42,11 +42,13 @@ class TurboQuantHadamardProgram final : public Program= uniforms.num_slices_per_kv) { return; } - // Compute total_seq_length + // Must match the batch-wide copy length used to build the host dispatch layout. + let copy_seq_length = uniforms.copy_sequence_length; + let batch = kv_slice / (uniforms.kv_num_heads * copy_seq_length); + let head = (kv_slice / copy_seq_length) % uniforms.kv_num_heads; + let seq = kv_slice % copy_seq_length; + + // Compute the logical total sequence length for this batch. #if use_seqlen_k - let total_seq_length = u32(seqlen_k[0u]) + 1u; + let per_batch_total_seq_length = u32(seqlen_k[batch]) + 1u; #else - let total_seq_length = uniforms.total_sequence_length; + let per_batch_total_seq_length = uniforms.total_sequence_length; #endif // uniforms.kv_sequence_length is the sequence length of the new key/values. - let past_seq_length = total_seq_length - uniforms.kv_sequence_length; - - // Unflatten kv_slice into (batch, head, seq). The seq extent depends on whether - // past tokens also need processing (copy) or just new tokens. -#if past_present_share_buffer - let copy_seq_length = uniforms.kv_sequence_length; -#elif has_past - let copy_seq_length = total_seq_length; -#else - let copy_seq_length = uniforms.kv_sequence_length; -#endif - let batch = kv_slice / (uniforms.kv_num_heads * copy_seq_length); - let head = (kv_slice / copy_seq_length) % uniforms.kv_num_heads; - let seq = kv_slice % copy_seq_length; + let past_seq_length = per_batch_total_seq_length - min(per_batch_total_seq_length, uniforms.kv_sequence_length); // Prepare indirect dispatch buffer (first workgroup, first thread only). #if prepare_indirect_dispatch if (workgroup_idx == 0u && local_idx == 0u) { - let num_total_seq_length_tile = (total_seq_length + uniforms.tile_size - 1u) / uniforms.tile_size; + let global_total_seq_length = u32(total_sequence_length_input[0]); + let num_total_seq_length_tile = + (global_total_seq_length + uniforms.tile_size - 1u) / uniforms.tile_size; populate_indirect_dispatch_buffer(num_total_seq_length_tile, uniforms.num_heads * uniforms.num_q_tiles, uniforms.batch_size); } #endif + // Skip slices beyond this batch's logical total sequence length. The host dispatch + // uses a uniform copy length, so shorter right-padded batches still receive workgroups. + if (seq >= per_batch_total_seq_length) { + return; + } + // Compute destination offset in present_key/present_value (u32 packed, always BNSH layout). #if past_present_share_buffer let dest_seq = past_seq_length + seq; @@ -86,7 +87,8 @@ $MAIN { // Handle past tokens: simple u32 word copy (already quantized). #if has_past if (seq < past_seq_length) { - let past_base = ((batch * uniforms.kv_num_heads + head) * past_seq_length + seq) * COMPRESSED_HEAD_U32; + let past_base = + ((batch * uniforms.kv_num_heads + head) * uniforms.past_input_seq_length + seq) * COMPRESSED_HEAD_U32; for (var i = local_idx; i < COMPRESSED_HEAD_U32; i += workgroup_size_x) { if (!is_value) { present_key.setByOffset(present_base + i, past_key.getByOffset(past_base + i)); diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 75c3e95970b2a..930c588d0fde9 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -3345,23 +3346,408 @@ TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefillNonFlashAttention_W #ifdef USE_WEBGPU // --------------------------------------------------------------------------- -// TurboQuant KV cache quantization tests. -// Tests exercise the TQ4 code paths in GroupQueryAttention + FlashAttention. +// WebGPU graph-capture and TurboQuant KV cache quantization tests. +// Tests exercise static-cache preprocessing and the TQ4 code paths in +// GroupQueryAttention + FlashAttention. // The helpers below reference webgpu::options::* constants, which are only // available when USE_WEBGPU is defined; guard the whole section so non-WebGPU // test builds (CPU/CUDA) still compile the rest of this file. // --------------------------------------------------------------------------- -// Helper: creates a WebGPU EP with TurboQuant 4-bit enabled. -static std::unique_ptr WebGpuEPWithTurboQuant4() { +static std::unique_ptr WebGpuEPForGqaOptions(bool enable_graph_capture, + bool enable_turbo_quant, + uint32_t multi_rotary_cache_concat_offset = 0) { ConfigOptions config_options{}; ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kStorageBufferCacheMode, webgpu::options::kBufferCacheMode_Disabled)); - ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kKvCacheQuantizationBits, - webgpu::options::kKvCacheQuantizationBits_4Bit)); + if (enable_turbo_quant) { + ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kKvCacheQuantizationBits, + webgpu::options::kKvCacheQuantizationBits_4Bit)); + } + if (enable_graph_capture) { + ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kEnableGraphCapture, + webgpu::options::kEnableGraphCapture_ON)); + } + if (multi_rotary_cache_concat_offset > 0) { + ORT_THROW_IF_ERROR(config_options.AddConfigEntry( + webgpu::options::kMultiRotaryCacheConcatOffset, + std::to_string(multi_rotary_cache_concat_offset).c_str())); + } return WebGpuExecutionProviderWithOptions(config_options); } +// Helper: creates a WebGPU EP with TurboQuant 4-bit enabled. +static std::unique_ptr WebGpuEPWithTurboQuant4(bool enable_graph_capture = false) { + return WebGpuEPForGqaOptions(enable_graph_capture, /*enable_turbo_quant=*/true); +} + +// Graph capture requires the indirect-dispatch dimensions to be prepared on the GPU. +// Verify that static-cache preprocessing uses the batch-wide total_sequence_length input +// instead of deriving the dispatch width from batch 0's (possibly shorter) seqlens_k value. The +// four-token input also makes batch 0's logical total shorter than kv_sequence_length, +// covering the right-padding underflow clamp with true static-cache aliasing. +static void RunIndirectDispatchGraphCapture(bool do_rotary, + bool enable_turbo_quant, + bool enable_multi_rotary_cache) { + constexpr int batch_size = 2; + constexpr int sequence_length = 4; + constexpr int short_total_sequence_length = 2; + constexpr int cache_sequence_length = 130; // Three 64-token attention tiles. + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 128; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + constexpr int packed_hidden_size = hidden_size + 2 * kv_hidden_size; + constexpr int compressed_head_size = head_size / 8 + 1; + constexpr uint32_t multi_rotary_cache_concat_offset = 4; + const int cache_head_size = enable_turbo_quant ? compressed_head_size : head_size; + + std::unique_ptr model; + { + std::unordered_map domain_to_version; + domain_to_version[kOnnxDomain] = 17; + domain_to_version[kMSDomain] = 1; + model = std::make_unique( + do_rotary ? "tq_gc_rotary_test" : "tq_gc_test", true, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, + std::vector{}, + DefaultLoggingManager().DefaultLogger(), + ModelOptions(true, true)); + onnxruntime::Graph& graph = model->MainGraph(); + + ONNX_NAMESPACE::TypeProto tp_float, tp_int32; + tp_float.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tp_int32.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_INT32); + + auto& empty = graph.GetOrCreateNodeArg("", nullptr); + std::vector inputs = { + &graph.GetOrCreateNodeArg("query", &tp_float), + do_rotary ? &empty : &graph.GetOrCreateNodeArg("key", &tp_float), + do_rotary ? &empty : &graph.GetOrCreateNodeArg("value", &tp_float), + &graph.GetOrCreateNodeArg("past_key", &tp_float), + &graph.GetOrCreateNodeArg("past_value", &tp_float), + &graph.GetOrCreateNodeArg("seqlens_k", &tp_int32), + &graph.GetOrCreateNodeArg("total_sequence_length", &tp_int32), + do_rotary ? &graph.GetOrCreateNodeArg("cos_cache", &tp_float) : &empty, + do_rotary ? &graph.GetOrCreateNodeArg("sin_cache", &tp_float) : &empty, + &empty, // position_ids + &empty, // attention_bias + &empty, // head_sink + }; + std::vector outputs = { + &graph.GetOrCreateNodeArg("output", &tp_float), + &graph.GetOrCreateNodeArg("present_key", &tp_float), + &graph.GetOrCreateNodeArg("present_value", &tp_float), + }; + + auto& node = graph.AddNode("gqa", "GroupQueryAttention", "GQA", inputs, outputs, nullptr, kMSDomain); + node.AddAttribute("num_heads", static_cast(num_heads)); + node.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + if (do_rotary) { + node.AddAttribute("do_rotary", int64_t{1}); + } + ORT_THROW_IF_ERROR(graph.Resolve()); + } + + std::string model_data; + model->ToProto().SerializeToString(&model_data); + + SessionOptions session_options; + InferenceSession session{session_options, GetEnvironment()}; + auto webgpu_ep = WebGpuEPForGqaOptions( + /*enable_graph_capture=*/true, + enable_turbo_quant, + enable_multi_rotary_cache ? multi_rotary_cache_concat_offset : 0); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU EP not available"; + } + IExecutionProvider* ep = webgpu_ep.get(); + ORT_THROW_IF_ERROR(session.RegisterExecutionProvider(std::move(webgpu_ep))); + std::istringstream model_stream(model_data); + ORT_THROW_IF_ERROR(session.Load(model_stream)); + ORT_THROW_IF_ERROR(session.Initialize()); + + OrtMemoryInfo gpu_memory_info(WEBGPU_BUFFER, OrtAllocatorType::OrtDeviceAllocator, + OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, + OrtDevice::VendorIds::NONE, 0)); + auto gpu_allocator = session.GetAllocator(gpu_memory_info); + AllocatorPtr cpu_allocator = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; + + auto make_data = [](size_t size, float scale, int period) { + std::vector data(size); + for (size_t i = 0; i < size; ++i) { + data[i] = scale * static_cast(i % static_cast(period) + 1); + } + return data; + }; + auto swap_batches = [](const std::vector& input) { + const size_t elements_per_batch = input.size() / batch_size; + std::vector output(input.size()); + std::copy_n(input.data() + elements_per_batch, elements_per_batch, output.data()); + std::copy_n(input.data(), elements_per_batch, output.data() + elements_per_batch); + return output; + }; + + const int query_width = do_rotary ? packed_hidden_size : hidden_size; + auto query_data = make_data(batch_size * sequence_length * query_width, 0.01f, 31); + auto key_data = make_data(batch_size * sequence_length * kv_hidden_size, 0.02f, 29); + auto value_data = make_data(batch_size * sequence_length * kv_hidden_size, 0.03f, 23); + auto past_key_data = make_data(batch_size * kv_num_heads * cache_sequence_length * cache_head_size, + 0.001f, 19); + auto past_value_data = make_data(batch_size * kv_num_heads * cache_sequence_length * cache_head_size, + 0.002f, 17); + auto query_data_swapped = swap_batches(query_data); + auto key_data_swapped = swap_batches(key_data); + auto value_data_swapped = swap_batches(value_data); + auto past_key_data_swapped = swap_batches(past_key_data); + auto past_value_data_swapped = swap_batches(past_value_data); + + constexpr int half_rotary_dim = head_size / 2; + const int large_rotary_cache_length = cache_sequence_length + 1; + auto cos_cache_data = make_data(large_rotary_cache_length * half_rotary_dim, 0.001f, 37); + auto sin_cache_data = make_data(large_rotary_cache_length * half_rotary_dim, 0.001f, 41); + int rotary_cache_length = large_rotary_cache_length; + if (enable_multi_rotary_cache) { + const size_t small_cache_size = multi_rotary_cache_concat_offset * half_rotary_dim; + cos_cache_data.insert(cos_cache_data.begin(), small_cache_size, + std::numeric_limits::quiet_NaN()); + sin_cache_data.insert(sin_cache_data.begin(), small_cache_size, + std::numeric_limits::quiet_NaN()); + rotary_cache_length += multi_rotary_cache_concat_offset; + } + + auto make_gpu_value = [&](const void* data, MLDataType data_type, const TensorShape& shape) { + Tensor gpu_tensor(data_type, shape, gpu_allocator); + Tensor cpu_tensor(data_type, shape, const_cast(data), cpu_allocator->Info()); + ORT_THROW_IF_ERROR(ep->GetDataTransfer()->CopyTensor(cpu_tensor, gpu_tensor)); + OrtValue value; + Tensor::InitOrtValue(std::move(gpu_tensor), value); + return value; + }; + auto update_gpu_value = [&](const OrtValue& gpu_value, const void* data, MLDataType data_type, + const TensorShape& shape) { + Tensor cpu_tensor(data_type, shape, const_cast(data), cpu_allocator->Info()); + ORT_THROW_IF_ERROR(ep->GetDataTransfer()->CopyTensor( + cpu_tensor, const_cast(gpu_value.Get()))); + }; + + const TensorShape query_shape{batch_size, sequence_length, query_width}; + const TensorShape kv_shape{batch_size, sequence_length, kv_hidden_size}; + const TensorShape cache_shape{batch_size, kv_num_heads, cache_sequence_length, cache_head_size}; + const TensorShape seqlens_shape{batch_size}; + const TensorShape total_sequence_length_shape{1}; + const TensorShape rotary_cache_shape{rotary_cache_length, half_rotary_dim}; + auto query_value = make_gpu_value(query_data.data(), DataTypeImpl::GetType(), query_shape); + auto key_value = make_gpu_value(key_data.data(), DataTypeImpl::GetType(), kv_shape); + auto value_value = make_gpu_value(value_data.data(), DataTypeImpl::GetType(), kv_shape); + auto past_key_value = make_gpu_value(past_key_data.data(), DataTypeImpl::GetType(), cache_shape); + auto past_value_value = make_gpu_value(past_value_data.data(), DataTypeImpl::GetType(), cache_shape); + std::vector seqlens_data{short_total_sequence_length - 1, cache_sequence_length - 1}; + auto seqlens_value = make_gpu_value(seqlens_data.data(), DataTypeImpl::GetType(), seqlens_shape); + std::vector total_sequence_length_data{cache_sequence_length}; + auto total_sequence_length_value = make_gpu_value(total_sequence_length_data.data(), + DataTypeImpl::GetType(), + total_sequence_length_shape); + auto cos_cache_value = make_gpu_value(cos_cache_data.data(), DataTypeImpl::GetType(), rotary_cache_shape); + auto sin_cache_value = make_gpu_value(sin_cache_data.data(), DataTypeImpl::GetType(), rotary_cache_shape); + + Tensor output_tensor(DataTypeImpl::GetType(), + TensorShape{batch_size, sequence_length, hidden_size}, gpu_allocator); + OrtValue output_value; + Tensor::InitOrtValue(std::move(output_tensor), output_value); + + std::unique_ptr io_binding; + ORT_THROW_IF_ERROR(session.NewIOBinding(&io_binding)); + ORT_THROW_IF_ERROR(io_binding->BindInput("query", query_value)); + if (!do_rotary) { + ORT_THROW_IF_ERROR(io_binding->BindInput("key", key_value)); + ORT_THROW_IF_ERROR(io_binding->BindInput("value", value_value)); + } + ORT_THROW_IF_ERROR(io_binding->BindInput("past_key", past_key_value)); + ORT_THROW_IF_ERROR(io_binding->BindInput("past_value", past_value_value)); + ORT_THROW_IF_ERROR(io_binding->BindInput("seqlens_k", seqlens_value)); + ORT_THROW_IF_ERROR(io_binding->BindInput("total_sequence_length", total_sequence_length_value)); + if (do_rotary) { + ORT_THROW_IF_ERROR(io_binding->BindInput("cos_cache", cos_cache_value)); + ORT_THROW_IF_ERROR(io_binding->BindInput("sin_cache", sin_cache_value)); + } + ORT_THROW_IF_ERROR(io_binding->BindOutput("output", output_value)); + // Alias past and present buffers so packed rotary uses the static-cache fused path. + ORT_THROW_IF_ERROR(io_binding->BindOutput("present_key", past_key_value)); + ORT_THROW_IF_ERROR(io_binding->BindOutput("present_value", past_value_value)); + ORT_THROW_IF_ERROR(io_binding->SynchronizeInputs()); + + auto read_output = [&]() { + auto& gpu_output = io_binding->GetOutputs()[0].Get(); + Tensor cpu_output(DataTypeImpl::GetType(), gpu_output.Shape(), cpu_allocator); + ORT_THROW_IF_ERROR(ep->GetDataTransfer()->CopyTensor(gpu_output, cpu_output)); + return std::vector(cpu_output.Data(), cpu_output.Data() + cpu_output.Shape().Size()); + }; + auto read_gpu_bytes = [&](const OrtValue& gpu_value) { + const auto& gpu_tensor = gpu_value.Get(); + Tensor cpu_tensor(gpu_tensor.DataType(), gpu_tensor.Shape(), cpu_allocator); + ORT_THROW_IF_ERROR(ep->GetDataTransfer()->CopyTensor(gpu_tensor, cpu_tensor)); + const auto* begin = static_cast(cpu_tensor.DataRaw()); + return std::vector(begin, begin + cpu_tensor.SizeInBytes()); + }; + + RunOptions run_options; + ORT_THROW_IF_ERROR(session.Run(run_options, *io_binding)); + auto first_output = read_output(); + + // Batch 0 has only two logical tokens in a four-token input. TurboQuant static-cache + // slots for its two padded tokens must retain their original contents. The standard + // path currently writes padding slots, which is unrelated to cache-bank selection. + auto expect_padding_unchanged = [&](const std::vector& actual, + const std::vector& initial, + const char* cache_name) { + ASSERT_EQ(actual.size(), initial.size() * sizeof(float)); + const size_t bytes_per_token = kv_num_heads * cache_head_size * sizeof(float); + const size_t padding_offset = short_total_sequence_length * bytes_per_token; + const size_t padding_size = (sequence_length - short_total_sequence_length) * bytes_per_token; + const auto* initial_bytes = reinterpret_cast(initial.data()); + EXPECT_TRUE(std::equal(actual.begin() + padding_offset, + actual.begin() + padding_offset + padding_size, + initial_bytes + padding_offset)) + << cache_name << " padded static-cache slots were overwritten"; + }; + if (enable_turbo_quant) { + expect_padding_unchanged(read_gpu_bytes(past_key_value), past_key_data, "key"); + expect_padding_unchanged(read_gpu_bytes(past_value_value), past_value_data, "value"); + } + + update_gpu_value(query_value, query_data_swapped.data(), DataTypeImpl::GetType(), query_shape); + if (!do_rotary) { + update_gpu_value(key_value, key_data_swapped.data(), DataTypeImpl::GetType(), kv_shape); + update_gpu_value(value_value, value_data_swapped.data(), DataTypeImpl::GetType(), kv_shape); + } + update_gpu_value(past_key_value, past_key_data_swapped.data(), DataTypeImpl::GetType(), cache_shape); + update_gpu_value(past_value_value, past_value_data_swapped.data(), DataTypeImpl::GetType(), cache_shape); + seqlens_data = {cache_sequence_length - 1, short_total_sequence_length - 1}; + update_gpu_value(seqlens_value, seqlens_data.data(), DataTypeImpl::GetType(), seqlens_shape); + ORT_THROW_IF_ERROR(session.Run(run_options, *io_binding)); + auto second_output = read_output(); + + ASSERT_EQ(first_output.size(), second_output.size()); + EXPECT_TRUE(std::all_of(first_output.begin(), first_output.end(), + [](float value) { return std::isfinite(value); })) + << "first graph-capture output contains a non-finite value"; + EXPECT_TRUE(std::all_of(second_output.begin(), second_output.end(), + [](float value) { return std::isfinite(value); })) + << "second graph-capture output contains a non-finite value"; + constexpr size_t output_elements_per_batch = sequence_length * hidden_size; + for (int second_batch = 0; second_batch < batch_size; ++second_batch) { + const int first_batch = batch_size - 1 - second_batch; + const auto* second_begin = second_output.data() + second_batch * output_elements_per_batch; + const auto* first_begin = first_output.data() + first_batch * output_elements_per_batch; + EXPECT_TRUE(std::equal(second_begin, second_begin + output_elements_per_batch, first_begin)) + << "indirect-dispatch output mismatch after swapping batch " << first_batch + << " into slot " << second_batch; + } +} + +TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_IndirectDispatch_UsesGlobalLength_NoRotary) { + RunIndirectDispatchGraphCapture(/*do_rotary=*/false, + /*enable_turbo_quant=*/true, + /*enable_multi_rotary_cache=*/false); +} + +TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_IndirectDispatch_UsesGlobalLength_Rotary) { + RunIndirectDispatchGraphCapture(/*do_rotary=*/true, + /*enable_turbo_quant=*/true, + /*enable_multi_rotary_cache=*/false); +} + +TEST(GroupQueryAttentionTest, WebGPU_IndirectDispatch_MultiRotaryCache_UsesGlobalLength) { + RunIndirectDispatchGraphCapture(/*do_rotary=*/true, + /*enable_turbo_quant=*/false, + /*enable_multi_rotary_cache=*/true); +} + +TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_IndirectDispatch_MultiRotaryCache_UsesGlobalLength) { + RunIndirectDispatchGraphCapture(/*do_rotary=*/true, + /*enable_turbo_quant=*/true, + /*enable_multi_rotary_cache=*/true); +} + +// The non-static packed-QKV path uses split_packed_qkv_with_rotary_embedding. +// A batch-wide total above the concat offset must select the long RoPE cache for +// every batch, including batches whose individual total remains below the offset. +TEST(GroupQueryAttentionTest, WebGPU_MultiRotaryCache_UsesGlobalLength_NonStaticCache) { + constexpr int batch_size = 2; + constexpr int sequence_length = 1; + constexpr int past_sequence_length = 4; + constexpr int total_sequence_length = past_sequence_length + sequence_length; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 16; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + constexpr int packed_hidden_size = hidden_size + 2 * kv_hidden_size; + constexpr int half_rotary_dim = head_size / 2; + constexpr uint32_t multi_rotary_cache_concat_offset = 4; + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", num_heads); + tester.AddAttribute("kv_num_heads", kv_num_heads); + tester.AddAttribute("do_rotary", 1); + + std::vector packed_qkv(batch_size * sequence_length * packed_hidden_size); + for (size_t i = 0; i < packed_qkv.size(); ++i) { + packed_qkv[i] = 0.01f * static_cast(i % 17 + 1); + } + tester.AddInput("query", {batch_size, sequence_length, packed_hidden_size}, packed_qkv); + tester.AddOptionalInputEdge(); // key + tester.AddOptionalInputEdge(); // value + + const int past_cache_size = batch_size * kv_num_heads * past_sequence_length * head_size; + tester.AddInput("past_key", {batch_size, kv_num_heads, past_sequence_length, head_size}, + std::vector(past_cache_size, 0.05f)); + tester.AddInput("past_value", {batch_size, kv_num_heads, past_sequence_length, head_size}, + std::vector(past_cache_size, 0.07f)); + + // Per-batch totals are {3, 5}; only the batch-global total crosses offset 4. + tester.AddInput("seqlens_k", {batch_size}, {2, 4}); + tester.AddInput("total_sequence_length", {1}, {total_sequence_length}, /*is_initializer=*/true); + + const int large_cache_length = total_sequence_length; + const size_t small_cache_size = multi_rotary_cache_concat_offset * half_rotary_dim; + std::vector cos_cache(small_cache_size, std::numeric_limits::quiet_NaN()); + std::vector sin_cache(small_cache_size, std::numeric_limits::quiet_NaN()); + cos_cache.insert(cos_cache.end(), large_cache_length * half_rotary_dim, 1.0f); + sin_cache.insert(sin_cache.end(), large_cache_length * half_rotary_dim, 0.0f); + tester.AddInput("cos_cache", {multi_rotary_cache_concat_offset + large_cache_length, half_rotary_dim}, + cos_cache); + tester.AddInput("sin_cache", {multi_rotary_cache_concat_offset + large_cache_length, half_rotary_dim}, + sin_cache); + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, + std::vector(batch_size * sequence_length * hidden_size, 0.0f)); + const int present_cache_size = batch_size * kv_num_heads * total_sequence_length * head_size; + tester.AddOutput("present_key", {batch_size, kv_num_heads, total_sequence_length, head_size}, + std::vector(present_cache_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, total_sequence_length, head_size}, + std::vector(present_cache_size, 0.0f)); + tester.SetCustomOutputVerifier([](const std::vector& fetches, const std::string&) { + const auto& output = fetches[0].Get(); + const float* output_data = output.Data(); + EXPECT_TRUE(std::all_of(output_data, output_data + output.Shape().Size(), + [](float value) { return std::isfinite(value); })) + << "multi-RoPE output contains a non-finite value"; + }); + + std::vector> execution_providers; + execution_providers.push_back(WebGpuEPForGqaOptions( + /*enable_graph_capture=*/false, + /*enable_turbo_quant=*/false, + multi_rotary_cache_concat_offset)); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + // Helper to run a GQA op with TurboQuant enabled and separate Q/K/V with rotary. // past_seq_len controls total KV cache depth; sequence_length controls prefill vs decode. // Returns the output tensor data on success. @@ -3746,23 +4132,301 @@ TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_Prefill_PackedRotary_K24) { EXPECT_FALSE(all_zero) << "TurboQuant prefill packed+rotary K=24 output should not be all zeros"; } -// --- Error path: multi-batch with per-batch seqlens_k is rejected with TurboQuant --- -// The TurboQuant copy-to-quantized-KV-cache kernel reads seqlen_k[0] for every -// batch on the graph-capture decode path, so it only supports batch_size == 1 and -// explicitly rejects batch_size > 1 rather than silently corrupting batches 1..N-1. -// (The non-quantized flash-attention copy path does support per-batch seqlens_k, so -// this restriction is specific to KV cache quantization.) genai decode runs -// batch_size==1, so multi-batch is not a supported production path. -TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_RejectsMultiBatch) { - auto ep = WebGpuEPWithTurboQuant4(); - if (!ep) { +// --- Decode test helper: multi-batch with per-batch seqlens_k using TurboQuant --- +// Before the fix, the TurboQuant copy-to-quantized-KV-cache kernels read seqlen_k[0] +// for EVERY batch, so batches 1..N-1 used the wrong past length. This helper proves the +// kernels now read seqlen_k[batch] via "swap invariance": +// +// Run 1: batch data [A, B], seqlens_k [sA, sB] +// Run 2: batch data [B, A], seqlens_k [sB, sA] (both batches physically swapped) +// +// If each batch correctly uses its own seqlen, then the valid cache prefixes swap exactly +// between runs. If the kernel wrongly used seqlen_k[0] for all batches, the two runs would +// copy from or write to different cache positions for the same data and the swapped caches +// would NOT match. +// +// Both variants exercise turbo_quant_hadamard. The rotary variant additionally covers +// the separate Q/K rotary preprocessing used when past/present buffers are not aliased. +static void RunTurboQuantMultiBatchSwapInvariance(bool do_rotary) { + if (!WebGpuEPWithTurboQuant4()) { GTEST_SKIP() << "WebGPU EP not available"; } - RunGQATurboQuant(/*batch_size=*/2, /*sequence_length=*/1, /*past_seq_len=*/24, - /*num_heads=*/2, /*kv_num_heads=*/1, /*head_size=*/128, - /*do_rotary=*/false, /*is_packed_qkv=*/false, - OpTester::ExpectResult::kExpectFailure, - "supports batch_size == 1 only"); + + constexpr int batch_size = 2; + constexpr int sequence_length = 1; + constexpr int past_seq_len = 24; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 128; + + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + constexpr int total_sequence_length = past_seq_len + sequence_length; + constexpr int kv_head_dim = (head_size * 4 + 32) / 32; // TQ4 compressed dim + constexpr int output_size = batch_size * sequence_length * hidden_size; + constexpr int present_size = batch_size * kv_num_heads * total_sequence_length * kv_head_dim; + + // Distinct per-batch past lengths (right-padded prompts of different lengths). + constexpr int32_t seqA = 20; // batch slot 0 "A" + constexpr int32_t seqB = 24; // batch slot 1 "B" + + std::mt19937 rng(42); + std::uniform_real_distribution dist(-0.5f, 0.5f); + + // Per-batch Q/K/V and (already-packed) past-KV blocks for the two logical prompts A and B. + const int q_per_batch = sequence_length * hidden_size; + const int kv_per_batch = sequence_length * kv_hidden_size; + const int past_per_batch = kv_num_heads * past_seq_len * kv_head_dim; + + auto make_vec = [&](int n) { + std::vector v(n); + for (auto& e : v) e = dist(rng); + return v; + }; + const std::vector qA = make_vec(q_per_batch), qB = make_vec(q_per_batch); + const std::vector kA = make_vec(kv_per_batch), kB = make_vec(kv_per_batch); + const std::vector vA = make_vec(kv_per_batch), vB = make_vec(kv_per_batch); + const std::vector pkA = make_vec(past_per_batch), pkB = make_vec(past_per_batch); + const std::vector pvA = make_vec(past_per_batch), pvB = make_vec(past_per_batch); + + const int max_seq_len = total_sequence_length + 8; + const int half_rotary = head_size / 2; + std::vector cos_cache(max_seq_len * half_rotary); + std::vector sin_cache(max_seq_len * half_rotary); + for (int pos = 0; pos < max_seq_len; ++pos) { + for (int d = 0; d < half_rotary; ++d) { + float freq = 1.0f / std::pow(10000.0f, 2.0f * static_cast(d) / static_cast(head_size)); + cos_cache[pos * half_rotary + d] = std::cos(static_cast(pos) * freq); + sin_cache[pos * half_rotary + d] = std::sin(static_cast(pos) * freq); + } + } + + auto concat = [](const std::vector& a, const std::vector& b) { + std::vector out(a); + out.insert(out.end(), b.begin(), b.end()); + return out; + }; + + struct PackedPresentCache { + std::vector key; + std::vector value; + }; + + auto copy_tensor_bytes = [](const Tensor& tensor) { + const auto* begin = static_cast(tensor.DataRaw()); + return std::vector(begin, begin + tensor.SizeInBytes()); + }; + + // Runs GQA with the two batch slots holding (first, second) prompts and the given seqlens. + // Returns owned copies of the packed present caches. + auto run = [&](const std::vector& q0, const std::vector& q1, + const std::vector& k0, const std::vector& k1, + const std::vector& v0, const std::vector& v1, + const std::vector& pk0, const std::vector& pk1, + const std::vector& pv0, const std::vector& pv1, + int32_t s0, int32_t s1) { + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + if (do_rotary) { + tester.AddAttribute("do_rotary", static_cast(1)); + } + + tester.AddInput("query", {batch_size, sequence_length, hidden_size}, concat(q0, q1)); + tester.AddInput("key", {batch_size, sequence_length, kv_hidden_size}, concat(k0, k1)); + tester.AddInput("value", {batch_size, sequence_length, kv_hidden_size}, concat(v0, v1)); + tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, kv_head_dim}, concat(pk0, pk1)); + tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, kv_head_dim}, concat(pv0, pv1)); + + std::vector seqlens_k{s0, s1}; + tester.AddInput("seqlens_k", {batch_size}, seqlens_k); + tester.AddInput("total_sequence_length", {1}, {total_sequence_length}, /*is_initializer=*/true); + + if (do_rotary) { + tester.AddInput("cos_cache", {max_seq_len, half_rotary}, cos_cache); + tester.AddInput("sin_cache", {max_seq_len, half_rotary}, sin_cache); + } else { + tester.AddOptionalInputEdge(); // cos_cache + tester.AddOptionalInputEdge(); // sin_cache + } + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, + std::vector(output_size, 0.0f)); + tester.AddOutput("present_key", {batch_size, kv_num_heads, total_sequence_length, kv_head_dim}, + std::vector(present_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, total_sequence_length, kv_head_dim}, + std::vector(present_size, 0.0f)); + + // Only the attention output is numerically meaningful; present_key/value are packed + // quantized bytes reinterpreted as float, so skip their value checks. + tester.SetOutputTolerance(1e6f); + tester.SetCustomOutputVerifier([](const std::vector&, const std::string&) {}); + + std::vector> execution_providers; + execution_providers.push_back(WebGpuEPWithTurboQuant4()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + + auto fetches = tester.GetFetches(); + return PackedPresentCache{copy_tensor_bytes(fetches[1].Get()), + copy_tensor_bytes(fetches[2].Get())}; + }; + + // Run 1: slot0=A (seqA), slot1=B (seqB). + auto out1 = run(qA, qB, kA, kB, vA, vB, pkA, pkB, pvA, pvB, seqA, seqB); + // Run 2: swap both the data AND the seqlens so slot0=B (seqB), slot1=A (seqA). + auto out2 = run(qB, qA, kB, kA, vB, vA, pkB, pkA, pvB, pvA, seqB, seqA); + + auto expect_batches_swapped = [&](const std::vector& first, + const std::vector& second, + const char* cache_name) { + ASSERT_EQ(first.size(), second.size()); + ASSERT_EQ(first.size(), static_cast(present_size) * sizeof(float)); + const size_t bytes_per_batch = first.size() / batch_size; + constexpr size_t bytes_per_token = kv_num_heads * kv_head_dim * sizeof(float); + for (int second_batch = 0; second_batch < batch_size; ++second_batch) { + const int first_batch = batch_size - 1 - second_batch; + const auto* second_begin = second.data() + second_batch * bytes_per_batch; + const auto* first_begin = first.data() + first_batch * bytes_per_batch; + const int32_t logical_sequence_length = first_batch == 0 ? seqA + 1 : seqB + 1; + const size_t valid_bytes = logical_sequence_length * bytes_per_token; + EXPECT_TRUE(std::equal(second_begin, second_begin + valid_bytes, first_begin)) + << cache_name << " cache mismatch after swapping batch " << first_batch + << " into slot " << second_batch + << "; per-batch seqlens_k not honored in quantized KV cache path"; + } + }; + + expect_batches_swapped(out1.key, out2.key, "key"); + expect_batches_swapped(out1.value, out2.value, "value"); +} + +// Rotary variant: exercises the fused rotary+Hadamard copy kernel. +TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_Decode_MultiBatch_UsesPerBatchSeqlensK) { + RunTurboQuantMultiBatchSwapInvariance(/*do_rotary=*/true); +} + +// Non-rotary variant: exercises the plain Hadamard copy kernel (turbo_quant_hadamard). +TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_Decode_MultiBatch_NoRotary_UsesPerBatchSeqlensK) { + RunTurboQuantMultiBatchSwapInvariance(/*do_rotary=*/false); +} + +// Right-padded first prompts can have a per-batch total sequence length smaller than the +// padded K/V sequence length. Exercise the dynamic-cache path for turbo_quant_hadamard; +// the graph-capture tests above cover the static-cache and fused rotary variants. +static void RunTurboQuantRightPaddedPrefill(bool do_rotary) { + if (!WebGpuEPWithTurboQuant4()) { + GTEST_SKIP() << "WebGPU EP not available"; + } + + constexpr int batch_size = 2; + constexpr int sequence_length = 4; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 128; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + constexpr int packed_hidden_size = hidden_size + 2 * kv_hidden_size; + constexpr int kv_head_dim = (head_size * 4 + 32) / 32; + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", num_heads); + tester.AddAttribute("kv_num_heads", kv_num_heads); + if (do_rotary) { + tester.AddAttribute("do_rotary", 1); + } + + std::mt19937 rng(2026); + std::uniform_real_distribution dist(-0.5f, 0.5f); + auto make_data = [&](int size) { + std::vector data(size); + for (float& value : data) { + value = dist(rng); + } + return data; + }; + + if (do_rotary) { + tester.AddInput("query", {batch_size, sequence_length, packed_hidden_size}, + make_data(batch_size * sequence_length * packed_hidden_size)); + tester.AddOptionalInputEdge(); // key + tester.AddOptionalInputEdge(); // value + } else { + tester.AddInput("query", {batch_size, sequence_length, hidden_size}, + make_data(batch_size * sequence_length * hidden_size)); + tester.AddInput("key", {batch_size, sequence_length, kv_hidden_size}, + make_data(batch_size * sequence_length * kv_hidden_size)); + tester.AddInput("value", {batch_size, sequence_length, kv_hidden_size}, + make_data(batch_size * sequence_length * kv_hidden_size)); + } + + const int cache_size = batch_size * kv_num_heads * sequence_length * kv_head_dim; + tester.AddInput("past_key", {batch_size, kv_num_heads, sequence_length, kv_head_dim}, + std::vector(cache_size, 0.0f)); + tester.AddInput("past_value", {batch_size, kv_num_heads, sequence_length, kv_head_dim}, + std::vector(cache_size, 0.0f)); + + // seqlens_k is total_sequence_length - 1. Batch 0 therefore has only two valid + // tokens in the four-token padded prompt, while batch 1 uses the full prompt. + tester.AddInput("seqlens_k", {batch_size}, {1, 3}); + tester.AddInput("total_sequence_length", {1}, {sequence_length}, /*is_initializer=*/true); + + if (do_rotary) { + constexpr int max_seq_len = sequence_length + 8; + constexpr int half_rotary = head_size / 2; + std::vector cos_cache(max_seq_len * half_rotary); + std::vector sin_cache(max_seq_len * half_rotary); + for (int pos = 0; pos < max_seq_len; ++pos) { + for (int dim = 0; dim < half_rotary; ++dim) { + const float frequency = 1.0f / std::pow(10000.0f, 2.0f * static_cast(dim) / head_size); + cos_cache[pos * half_rotary + dim] = std::cos(static_cast(pos) * frequency); + sin_cache[pos * half_rotary + dim] = std::sin(static_cast(pos) * frequency); + } + } + tester.AddInput("cos_cache", {max_seq_len, half_rotary}, cos_cache); + tester.AddInput("sin_cache", {max_seq_len, half_rotary}, sin_cache); + } else { + tester.AddOptionalInputEdge(); // cos_cache + tester.AddOptionalInputEdge(); // sin_cache + } + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + const int output_size = batch_size * sequence_length * hidden_size; + tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, + std::vector(output_size, 0.0f)); + tester.AddOutput("present_key", {batch_size, kv_num_heads, sequence_length, kv_head_dim}, + std::vector(cache_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, sequence_length, kv_head_dim}, + std::vector(cache_size, 0.0f)); + tester.SetOutputTolerance(1e6f); + tester.SetCustomOutputVerifier([](const std::vector&, const std::string&) {}); + + std::vector> execution_providers; + execution_providers.push_back(WebGpuEPWithTurboQuant4()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + + auto fetches = tester.GetFetches(); + const auto& output = fetches[0].Get(); + const float* output_data = output.Data(); + const int output_per_batch = sequence_length * hidden_size; + for (int batch = 0; batch < batch_size; ++batch) { + const float* batch_begin = output_data + batch * output_per_batch; + const bool all_zero = std::all_of(batch_begin, batch_begin + output_per_batch, + [](float value) { return value == 0.0f; }); + EXPECT_FALSE(all_zero) << "TurboQuant output is all zero for right-padded batch " << batch; + } +} + +TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_Prefill_MultiBatch_RightPadding_NoRotary) { + RunTurboQuantRightPaddedPrefill(/*do_rotary=*/false); +} + +TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_Prefill_MultiBatch_RightPadding_Rotary) { + RunTurboQuantRightPaddedPrefill(/*do_rotary=*/true); } // ---------------------------------------------------------------------------