[WebGPU] Initial PagedAttention implementation (1/n)

onnxruntime/contrib_ops/cuda/bert/paged_attention.cc

@@ -11,7 +11,7 @@
 #include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
 #include "contrib_ops/cuda/bert/paged_attention_impl.h"
 #include "contrib_ops/cuda/bert/paged_attention.h"
-#include "contrib_ops/cuda/bert/paged_attention_helper.h"
+#include "contrib_ops/cpu/bert/paged_attention_helper.h"
 #include "contrib_ops/cuda/bert/flash_attention/flash_api.h"
 #include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h"
 #include "contrib_ops/cuda/bert/xqa/xqa_paged_loader.h"

onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc

@@ -263,6 +263,9 @@ Status FlashAttentionProgram::GenerateShaderCode(ShaderHelper& shader) const {
   if (use_seqlen_k_) {
     shader.AddInput("seqlens_k", ShaderUsage::None);
   }
+  if (use_seqlens_q_) {
+    shader.AddInput("seqlens_q", ShaderUsage::None);
+  }
   if (has_head_sink_) {
     shader.AddInput("head_sink", ShaderUsage::UseUniform);
   }
@@ -282,6 +285,7 @@ Status FlashAttentionProgram::GenerateShaderCode(ShaderHelper& shader) const {
                              WGSL_TEMPLATE_PARAMETER(qkv_num_heads, qkv_num_heads_),
                              WGSL_TEMPLATE_PARAMETER(turbo_quant, turbo_quant_),
                              WGSL_TEMPLATE_PARAMETER(use_seqlen_k, use_seqlen_k_),
+                             WGSL_TEMPLATE_PARAMETER(use_seqlens_q, use_seqlens_q_),
                              WGSL_TEMPLATE_PARAMETER(use_shm_path, use_shm_path_));
 }
 
@@ -292,6 +296,9 @@ Status FlashAttentionDecodeQKVProgram::GenerateShaderCode(ShaderHelper& shader)
   if (use_seqlen_k_) {
     shader.AddInput("seqlens_k", ShaderUsage::None);
   }
+  if (use_seqlens_q_) {
+    shader.AddInput("seqlens_q", ShaderUsage::None);
+  }
   if (use_indirect_dispatch_) {
     // Global max total sequence length across batches (from GQA input #6).
     // Used in indirect-dispatch mode for the workgroup_idx slicing so that
@@ -322,6 +329,7 @@ Status FlashAttentionDecodeQKVProgram::GenerateShaderCode(ShaderHelper& shader)
                              WGSL_TEMPLATE_PARAMETER(turbo_quant, turbo_quant_),
                              WGSL_TEMPLATE_PARAMETER(use_indirect_dispatch, use_indirect_dispatch_),
                              WGSL_TEMPLATE_PARAMETER(use_seqlen_k, use_seqlen_k_),
+                             WGSL_TEMPLATE_PARAMETER(use_seqlens_q, use_seqlens_q_),
                              WGSL_TEMPLATE_PARAMETER(v_head_size_vec, head_size_vec_),
                              WGSL_TEMPLATE_VARIABLE(metadata, metadata),
                              WGSL_TEMPLATE_VARIABLE(out_split_vx, out_split_vx),
@@ -334,7 +342,8 @@ Status ComputeFlashAttentionDecodeQKV(onnxruntime::webgpu::ComputeContext& conte
                                       const Tensor* attention_bias, Tensor* out_split_vx, Tensor* present_key, Tensor* present_value,
                                       Tensor* metadata, const Tensor* seqlen_k,
                                       const WebgpuAttentionParameters& parameters, const Tensor* indirect_buffer, uint32_t num_total_seq_length_tile, uint32_t num_present_sequence_length_tile, uint32_t tile_size, bool use_indirect_dispatch, uint32_t present_sequence_length, uint32_t m_tile, bool use_seqlen_k, const Tensor* total_seqlen,
-                                      bool turbo_quant, int compressed_head_size_u32) {
+                                      bool turbo_quant, int compressed_head_size_u32,
+                                      bool use_seqlens_q, const Tensor* seqlens_q) {
   const float alpha = parameters.scale_ == 0.0f ? 1.f / sqrt(static_cast<float>(parameters.head_size_))
                                                 : parameters.scale_;
 
@@ -347,13 +356,16 @@ Status ComputeFlashAttentionDecodeQKV(onnxruntime::webgpu::ComputeContext& conte
 
   bool q_BNSH = parameters.qkv_format_ == Q_K_V_BNSH;
   bool is_unidirectional = parameters.is_unidirectional_;
-  FlashAttentionDecodeQKVProgram program{"FlashAttentionDecodeQKV", has_attention_bias, tile_size, head_size_vec, use_indirect_dispatch, q_BNSH, is_unidirectional, m_tile, use_seqlen_k, turbo_quant, compressed_head_size_u32};
+  FlashAttentionDecodeQKVProgram program{"FlashAttentionDecodeQKV", has_attention_bias, tile_size, head_size_vec, use_indirect_dispatch, q_BNSH, is_unidirectional, m_tile, use_seqlen_k, turbo_quant, compressed_head_size_u32, use_seqlens_q};
   program.AddInputs({{Q, ProgramTensorMetadataDependency::TypeAndRank, components},
                      {present_key, ProgramTensorMetadataDependency::TypeAndRank, kv_cache_components},
                      {present_value, ProgramTensorMetadataDependency::TypeAndRank, kv_cache_components}});
   if (use_seqlen_k) {
     program.AddInput({seqlen_k, ProgramTensorMetadataDependency::None});
   }
+  if (use_seqlens_q) {
+    program.AddInput({seqlens_q, ProgramTensorMetadataDependency::None});
+  }
   if (use_indirect_dispatch) {
     program.AddInput({total_seqlen, ProgramTensorMetadataDependency::None});
   }
@@ -385,7 +397,7 @@ Status ComputeFlashAttentionDecodeQKV(onnxruntime::webgpu::ComputeContext& conte
   // for decode, 64 threads with 8 vec4 K tiles for prefill.
   const uint32_t workgroup_size = (m_tile == 1u) ? 128u : 64u;
   program.SetWorkgroupSize(workgroup_size)
-      .CacheHint(tile_size, head_size_vec, has_attention_bias, use_indirect_dispatch, q_BNSH, is_unidirectional, m_tile, use_seqlen_k, turbo_quant, compressed_head_size_u32)
+      .CacheHint(tile_size, head_size_vec, has_attention_bias, use_indirect_dispatch, q_BNSH, is_unidirectional, m_tile, use_seqlen_k, turbo_quant, compressed_head_size_u32, use_seqlens_q)
       .AddUniformVariables({{static_cast<uint32_t>(vectorized_head_size)},
                             {static_cast<uint32_t>(parameters.total_sequence_length_)},
                             {static_cast<float>(alpha)},
@@ -470,8 +482,9 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co
                            Tensor* output, const Tensor* past_key, Tensor* present_key, const Tensor* past_value, Tensor* present_value,
                            const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context, const Tensor* seqlen_k,
                            const Tensor* cos_cache, const Tensor* sin_cache, const Tensor* head_sink,
-                           const Tensor* total_seqlen) {
+                           const Tensor* total_seqlen, const Tensor* seqlens_q) {
   constexpr uint32_t tile_size = 64;
+  const bool use_seqlens_q = seqlens_q != nullptr;
 
   const bool turbo_quant_enabled = context.KvCacheQuantizationEnabled();
   if (turbo_quant_enabled && (parameters.head_size_ < 8 || (parameters.head_size_ & (parameters.head_size_ - 1)) != 0)) {
@@ -728,7 +741,8 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co
                                   use_seqlen_k,
                                   has_head_sink,
                                   turbo_quant_enabled,
-                                  compressed_head_size_u32};
+                                  compressed_head_size_u32,
+                                  use_seqlens_q};
     // When TQ is active, KV cache is u32-packed — use u32 tensor views for present_key/present_value.
     const Tensor* fa_present_key = turbo_quant_enabled ? tq_present_key : present_key;
     const Tensor* fa_present_value = turbo_quant_enabled ? tq_present_value : present_value;
@@ -741,6 +755,9 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co
     if (use_seqlen_k) {
       program.AddInputs({{seqlen_k, ProgramTensorMetadataDependency::None}});
     }
+    if (use_seqlens_q) {
+      program.AddInputs({{seqlens_q, ProgramTensorMetadataDependency::None}});
+    }
     if (has_head_sink) {
       program.AddInputs({{head_sink, ProgramTensorMetadataDependency::Type}});
     }
@@ -764,7 +781,7 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co
 
     program.SetDispatchGroupSize(parameters.batch_size_ * parameters.num_heads_ * num_seq_tile)
         .SetWorkgroupSize(prefill_tile_size)
-        .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_, parameters.is_unidirectional_, is_qualcomm, is_nvidia, is_apple, has_subgroups, q_BNSH, use_seqlen_k, has_head_sink, turbo_quant_enabled, compressed_head_size_u32, program.max_k_step())
+        .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_, parameters.is_unidirectional_, is_qualcomm, is_nvidia, is_apple, has_subgroups, q_BNSH, use_seqlen_k, has_head_sink, turbo_quant_enabled, compressed_head_size_u32, program.max_k_step(), use_seqlens_q)
         .AddUniformVariables({{static_cast<uint32_t>(parameters.sequence_length_)},
                               {static_cast<uint32_t>(parameters.total_sequence_length_)},
                               {static_cast<uint32_t>(present_sequence_length)},
@@ -800,7 +817,8 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co
                                                        parameters, indirect_buffer_ptr, num_total_seq_length_tile,
                                                        num_present_sequence_length_tile, tile_size, use_indirect_dispatch,
                                                        present_sequence_length, m_tile, use_seqlen_k, total_seqlen,
-                                                       turbo_quant_enabled, compressed_head_size_u32));
+                                                       turbo_quant_enabled, compressed_head_size_u32,
+                                                       use_seqlens_q, seqlens_q));
 
     ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeVxReduce(context, &out_split_vx, &metadata, attn_output, seqlen_k, parameters,
                                                             num_total_seq_length_tile,

onnxruntime/contrib_ops/webgpu/bert/flash_attention.h

@@ -101,7 +101,8 @@ class FlashAttentionProgram final : public Program<FlashAttentionProgram> {
                         bool use_seqlen_k = false,
                         bool has_head_sink = false,
                         bool turbo_quant = false,
-                        int compressed_head_size_u32 = 0)
+                        int compressed_head_size_u32 = 0,
+                        bool use_seqlens_q = false)
       : Program{kernel_name},
         has_attention_bias_(has_attention_bias),
         is_qualcomm_(is_qualcomm),
@@ -115,7 +116,8 @@ class FlashAttentionProgram final : public Program<FlashAttentionProgram> {
         use_seqlen_k_(use_seqlen_k),
         has_head_sink_(has_head_sink),
         turbo_quant_(turbo_quant),
-        compressed_head_size_u32_(compressed_head_size_u32) {
+        compressed_head_size_u32_(compressed_head_size_u32),
+        use_seqlens_q_(use_seqlens_q) {
     if (use_shm_path_) {
       // Use shared-memory loop-based path with dynamic max_k_step.
       // Compute max_k_step from workgroup shared memory budget: k_tile + v_tile = 2 * element_size * head_size * max_k_step
@@ -162,6 +164,10 @@ class FlashAttentionProgram final : public Program<FlashAttentionProgram> {
   int max_k_step_;
   bool turbo_quant_;
   int compressed_head_size_u32_;
+  // Per-batch new-Q-length path (LEFT-aligned Q). When set, the shader reads
+  // seqlens_q[b] and computes past_sequence_length_b = total_kv_b - q_len_b.
+  // When unset (default), callers keep the uniform-q_len clamp path unchanged.
+  bool use_seqlens_q_;
 };
 
 class FlashAttentionDecodeQKVProgram final : public Program<FlashAttentionDecodeQKVProgram> {
@@ -172,8 +178,9 @@ class FlashAttentionDecodeQKVProgram final : public Program<FlashAttentionDecode
                                  bool is_unidirectional = false,
                                  uint32_t m_tile = 1,
                                  bool use_seqlen_k = false,
-                                 bool turbo_quant = false, int compressed_head_size_u32 = 0)
-      : Program{kernel_name}, has_attention_bias_(has_attention_bias), tile_size_(tile_size), head_size_vec_(head_size_vec), use_indirect_dispatch_(use_indirect_dispatch), q_BNSH_(q_BNSH), is_unidirectional_(is_unidirectional), m_tile_(m_tile), use_seqlen_k_(use_seqlen_k), turbo_quant_(turbo_quant), compressed_head_size_u32_(compressed_head_size_u32) {
+                                 bool turbo_quant = false, int compressed_head_size_u32 = 0,
+                                 bool use_seqlens_q = false)
+      : Program{kernel_name}, has_attention_bias_(has_attention_bias), tile_size_(tile_size), head_size_vec_(head_size_vec), use_indirect_dispatch_(use_indirect_dispatch), q_BNSH_(q_BNSH), is_unidirectional_(is_unidirectional), m_tile_(m_tile), use_seqlen_k_(use_seqlen_k), turbo_quant_(turbo_quant), compressed_head_size_u32_(compressed_head_size_u32), use_seqlens_q_(use_seqlens_q) {
   }
 
   Status GenerateShaderCode(ShaderHelper& sh) const override;
@@ -202,6 +209,8 @@ class FlashAttentionDecodeQKVProgram final : public Program<FlashAttentionDecode
   bool use_seqlen_k_;
   bool turbo_quant_;
   int compressed_head_size_u32_;
+  // See FlashAttentionProgram::use_seqlens_q_ for semantics.
+  bool use_seqlens_q_;
 };
 
 class FlashAttentionDecodeVxReduceProgram final : public Program<FlashAttentionDecodeVxReduceProgram> {
@@ -228,11 +237,14 @@ class FlashAttentionDecodeVxReduceProgram final : public Program<FlashAttentionD
   bool use_seqlen_k_;
 };
 
+// seqlens_q (optional): int32[batch_size] of per-batch new-Q lengths. Enables
+// LEFT-aligned variable-q_len callers (e.g. PagedAttention). Uniform-q_len
+// callers pass nullptr and keep the pre-existing clamped path.
 Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, const Tensor* attention_bias,
                            Tensor* output, const Tensor* past_key, Tensor* present_key, const Tensor* past_value, Tensor* present_value,
                            const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context, const Tensor* seqlen_k = nullptr,
                            const Tensor* cos_cache = nullptr, const Tensor* sin_cache = nullptr, const Tensor* head_sink = nullptr,
-                           const Tensor* total_seqlen = nullptr);
+                           const Tensor* total_seqlen = nullptr, const Tensor* seqlens_q = nullptr);
 
 bool CanApplyFlashAttention(const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context);
 

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

@@ -11,6 +11,7 @@
 #param turbo_quant
 #param compressed_head_size_u32
 #param use_seqlen_k
+#param use_seqlens_q
 #param use_shm_path
 #param max_k_step_param
 
@@ -338,11 +339,20 @@ $MAIN {
   let total_sequence_length = get_total_sequence_length(batch_idx);
 
 #if is_unidirectional
-  // If attention is unidirectional, set the loop bound to enforce causal masking.
-  // Right-padded batches with prompt shorter than new_sequence_length would underflow u32; clamp to 0.
+  // Causal masking loop bound.
+#if use_seqlens_q
+  // LEFT-aligned per-batch Q length: past_len_b = total_kv_b - q_len_b (>= 0).
+  let q_len_b = u32(seqlens_q[batch_idx]);
+  let past_sequence_length = total_sequence_length - q_len_b;
+#else
+  // Uniform q_len path. total_kv < new_sequence_length can only happen for
+  // LEFT-aligned callers with past = 0 + short prompt (e.g. GQA right-padded
+  // prefill); clamping to 0 gives the correct past_len_b there. Variable-
+  // q_len callers with past > 0 must set use_seqlens_q.
   let past_sequence_length = select(total_sequence_length - uniforms.new_sequence_length,
                                     0u,
                                     total_sequence_length <= uniforms.new_sequence_length);
+#endif
   let max_causal_len_for_workgroup = past_sequence_length +
                                      (workgroup_idx % uniforms.num_seq_tile + 1) * workgroup_size_x;
   let loop_bound = min(total_sequence_length, max_causal_len_for_workgroup);

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

@@ -12,6 +12,7 @@
 #param turbo_quant
 #param use_indirect_dispatch
 #param use_seqlen_k
+#param use_seqlens_q
 #param v_head_size_vec
 
 #use .getByOffset .setByOffset
@@ -214,10 +215,15 @@ $MAIN {
 
   // Reduce inner_qk_values to tile_qk, apply attention bias and causal mask
 #if is_unidirectional
-  // Right-padded batches with prompt shorter than new_sequence_length would underflow u32; clamp to 0.
+#if use_seqlens_q
+  // LEFT-aligned per-batch Q length: past_len_b = total_kv_b - q_len_b.
+  let q_len_b = u32(seqlens_q[batch_idx]);
+  let past_sequence_length = total_sequence_length - q_len_b;
+#else
   let past_sequence_length = select(total_sequence_length - uniforms.new_sequence_length,
                                     0u,
                                     total_sequence_length <= uniforms.new_sequence_length);
+#endif
 #endif
   for (var m = 0u; m < m_tile && q_base + m < uniforms.new_sequence_length; m++) {
     let q_idx = q_base + m;