webgpu: Fix TurboQuant quantized KV cache for batch>1 with per-batch seqlens

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<uint32_t>(dispatch_size)},
       {static_cast<uint32_t>(params.batch_size_)},
       {num_q_tiles},
+      {static_cast<uint32_t>(params.total_sequence_length_)},
   });
 
   program.SetDispatchGroupSize((dispatch_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE);

onnxruntime/contrib_ops/webgpu/bert/flash_attention.h

@@ -37,7 +37,8 @@ class SplitPackedQKVWithRotaryEmbeddingAndCopyKVProgram final : public Program<S
       {"tile_size", ProgramUniformVariableDataType::Uint32},
       {"dispatch_size", ProgramUniformVariableDataType::Uint32},
       {"batch_size", ProgramUniformVariableDataType::Uint32},
-      {"num_q_tiles", ProgramUniformVariableDataType::Uint32});
+      {"num_q_tiles", ProgramUniformVariableDataType::Uint32},
+      {"total_sequence_length", ProgramUniformVariableDataType::Uint32});
 
  private:
   const bool interleaved_;

onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc

@@ -27,6 +27,9 @@ Status SplitPackedQKVWithRotaryEmbeddingProgram::GenerateShaderCode(ShaderHelper
   const auto& seqlens = sh.AddInput("seqlens", ShaderUsage::UseUniform);
   const auto& cos_cache = sh.AddInput("cos_cache", ShaderUsage::UseUniform);
   const auto& sin_cache = sh.AddInput("sin_cache", ShaderUsage::UseUniform);
+  if (use_total_sequence_length_input_) {
+    sh.AddInput("total_sequence_length_input", ShaderUsage::None);
+  }
 
   const auto& query = sh.AddOutput("query", ShaderUsage::UseUniform);
   const auto& key = sh.AddOutput("key", ShaderUsage::UseUniform);
@@ -36,6 +39,7 @@ Status SplitPackedQKVWithRotaryEmbeddingProgram::GenerateShaderCode(ShaderHelper
                              WGSL_TEMPLATE_PARAMETER(interleaved, interleaved_),
                              WGSL_TEMPLATE_PARAMETER(multi_rotary_cache_concat_offset, multi_rotary_cache_concat_offset_),
                              WGSL_TEMPLATE_PARAMETER(use_multi_rotary_cache_concat, multi_rotary_cache_concat_offset_ > 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<uint32_t>(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<uint32_t>(params.kv_num_heads_)},
           {static_cast<uint32_t>(head_size_vec)},
           {static_cast<uint32_t>(half_rotary_embedding_dim_vec)},
+          {static_cast<uint32_t>(params.total_sequence_length_)},
           {static_cast<uint32_t>(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;

onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h

@@ -16,10 +16,13 @@ using namespace onnxruntime::webgpu;
 
 class SplitPackedQKVWithRotaryEmbeddingProgram final : public Program<SplitPackedQKVWithRotaryEmbeddingProgram> {
  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<SplitPacke
       {"kv_num_heads", ProgramUniformVariableDataType::Uint32},
       {"head_size", ProgramUniformVariableDataType::Uint32},
       {"half_rotary_dim", ProgramUniformVariableDataType::Uint32},
+      {"total_sequence_length", ProgramUniformVariableDataType::Uint32},
       {"dispatch_size", ProgramUniformVariableDataType::Uint32});
 
  private:
   const bool interleaved_;
   const uint32_t multi_rotary_cache_concat_offset_;
+  const bool use_total_sequence_length_input_;
 };
 
 class GroupQueryAttention final : public WebGpuKernel {

onnxruntime/contrib_ops/webgpu/bert/split_packed_qkv_with_rotary_embedding.wgsl.template

@@ -1,5 +1,6 @@
 #param interleaved
 #param multi_rotary_cache_concat_offset
+#param use_total_sequence_length_input
 #param use_multi_rotary_cache_concat
 
 #use guardAgainstOutOfBoundsWorkgroupSizes
@@ -29,12 +30,17 @@ $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);
     let position_id = past_seqlen + seq_idx;
+#if use_total_sequence_length_input
+    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

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);
   }

onnxruntime/contrib_ops/webgpu/bert/turbo_quant_fused_rotary_hadamard.wgsl.template

@@ -42,42 +42,58 @@ var<workgroup> scale_reduction_buffer : array<f32, 64>;
 var<workgroup> index_buffer : array<u32, HEAD_SIZE>;
 
 $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
 
   // 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;