Vulkan: support attn_mask and scale in llama.custom_sdpa / sdpa_with_kv_cache

The LLM SDPA path rejected both arguments at load. Sliding window attention
(e.g. Gemma) needs them: an additive mask in place of the causal one, and a
scale other than 1/sqrt(head_dim).

As in the CPU op, the mask is a 2D [S, C] tensor over every key in the cache
and cannot be combined with is_causal. The QK shaders get _mask variants that
add it instead of applying the causal mask; softmax and AV are given C - S in
place of input_pos so they attend to all C keys, which needs no shader change.

Author: msluszniak

backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.glsl

@@ -18,6 +18,8 @@ $if IO_STORAGE == "buffer":
   #define ATTN_WEIGHTS_BUFFER
 $if K_CACHE_STORAGE == "buffer":
   #define K_CACHE_BUFFER
+$if HAS_MASK:
+  #define HAS_MASK
 
 #define Q_LAYOUT DHSB
 #define K_LAYOUT DHSB
@@ -40,10 +42,14 @@ layout(std430) buffer;
 ${layout_declare_tensor(B, "w", "t_attn_weights", DTYPE, IO_STORAGE, is_scalar_array=False)}
 ${layout_declare_tensor(B, "r", "t_q", DTYPE, IO_STORAGE, is_scalar_array=False)}
 ${layout_declare_tensor(B, "r", "t_k", DTYPE, K_CACHE_STORAGE, is_scalar_array=False)}
+$if HAS_MASK:
+  ${layout_declare_tensor(B, "r", "t_mask", DTYPE, IO_STORAGE, is_scalar_array=True)}
 
 ${layout_declare_ubo(B, "ivec4", "q_sizes")}
 ${layout_declare_ubo(B, "ivec4", "k_sizes")}
 ${layout_declare_ubo(B, "int", "input_pos")}
+$if HAS_MASK:
+  ${layout_declare_ubo(B, "ivec4", "mask_sizes")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
@@ -118,7 +124,12 @@ void main() {
   // If the tile is completely inside the mask region, then there is no need to
   // compute the output tile. All the elements in the output tile can be set to
   // negative infinity.
+#ifdef HAS_MASK
+  // The additive attn_mask replaces the causal mask, so no tile can be skipped.
+  const bool tile_in_mask_region = false;
+#else
   bool tile_in_mask_region = c > (input_pos + s + (TILE_M - 1));
+#endif
   if (tile_in_mask_region) {
     const VEC4_T negative_infinity_vec = VEC4_T(negative_infinity_val);
     set_out_tile_to_vec(out_tile, negative_infinity_vec);
@@ -171,13 +182,18 @@ void main() {
     out_tile = partial_sums[0];
     // Apply scale and mask if the tile was not entirely in the mask region
     if (!tile_in_mask_region) {
+#ifdef HAS_MASK
+      apply_scale(out_tile, inv_scale);
+      apply_attn_mask(out_tile, c, s);
+#else
       VEC4_T inv_scale_vec = VEC4_T(inv_scale);
       apply_scale_and_mask(
         out_tile,
         inv_scale_vec,
         input_pos,
         c,
         s);
+#endif
     }
 
     store_attn_weight_tile_with_checks(

backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_coop.yaml

@@ -11,6 +11,7 @@ sdpa_compute_attn_weights_coop:
     K_CACHE_STORAGE: texture3d
     TILE_K4: 1
     TILE_N4: 1
+    HAS_MASK: false
   generate_variant_forall:
     combination:
       parameter_names: [IO_STORAGE, K_CACHE_STORAGE]
@@ -23,3 +24,5 @@ sdpa_compute_attn_weights_coop:
       - VALUE: half
   shader_variants:
     - NAME: sdpa_compute_attn_weights_coop
+    - NAME: sdpa_compute_attn_weights_coop_mask
+      HAS_MASK: true

backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.glsl

@@ -36,6 +36,8 @@ $else:
 
 $if HAS_BIAS:
   #define HAS_BIAS
+$if HAS_MASK:
+  #define HAS_MASK
 
 #define TILE_M4 ${TILE_M4}
 #define TILE_K4 ${TILE_K4}
@@ -56,13 +58,17 @@ ${layout_declare_tensor(B, "r", "t_q", IN_DTYPE, IO_STORAGE, is_scalar_array=Fal
 ${layout_declare_tensor(B, "r", "t_k", IN_DTYPE, K_CACHE_STORAGE, is_scalar_array=False)}
 $if HAS_BIAS:
   ${layout_declare_tensor(B, "r", "t_bias", IN_DTYPE, IO_STORAGE, is_scalar_array=False)}
+$if HAS_MASK:
+  ${layout_declare_tensor(B, "r", "t_mask", IN_DTYPE, IO_STORAGE, is_scalar_array=True)}
 
 ${layout_declare_ubo(B, "ivec4", "q_sizes")}
 ${layout_declare_ubo(B, "ivec4", "k_sizes")}
 $if MODE == "llm":
   ${layout_declare_ubo(B, "int", "input_pos")}
 $if HAS_BIAS:
   ${layout_declare_ubo(B, "ivec4", "bias_sizes")}
+$if HAS_MASK:
+  ${layout_declare_ubo(B, "ivec4", "mask_sizes")}
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
@@ -82,6 +88,9 @@ ${layout_declare_spec_const(C, "float", "inv_scale", "1.0")}
  *   attn_weights: [B, H_q, S, context_len] in input dtype
  *   current context_len = input_pos + S
  *   Applies combined scale + causal mask.
+ *   With HAS_MASK the causal mask is replaced by an additive attn_mask
+ *   [S, context_len]; the caller passes input_pos = C - S so that every key
+ *   in k is attended, as in the CPU op.
  *
  * Fused SDPA:
  *   q:            [B, H, S, D]          (DSHB layout)
@@ -164,7 +173,7 @@ void main() {
   const int attn_S = S;
 #endif
 
-#ifdef HAS_INPUT_POS
+#if defined(HAS_INPUT_POS) && !defined(HAS_MASK)
   // If the tile is completely inside the mask region, then there is no need to
   // compute the output tile. All the elements in the output tile can be set to
   // negative infinity.
@@ -200,7 +209,11 @@ void main() {
       fp_accumulate_with_fp_weight(out_tile, q_tile, w_tile);
     }
 
-#ifdef HAS_INPUT_POS
+#if defined(HAS_INPUT_POS) && defined(HAS_MASK)
+    // LLM: scale + additive attn_mask
+    apply_scale(out_tile, inv_scale);
+    apply_attn_mask(out_tile, c, s);
+#elif defined(HAS_INPUT_POS)
     // LLM: combined scale + causal mask
     VEC4_T inv_scale_vec = VEC4_T(inv_scale);
     apply_scale_and_mask(

backends/vulkan/runtime/graph/ops/glsl/sdpa_compute_attn_weights_tiled.yaml

@@ -12,6 +12,7 @@ sdpa_compute_attn_weights_tiled:
     K_CACHE_STORAGE: texture3d
     MODE: llm
     HAS_BIAS: false
+    HAS_MASK: false
     TILE_M4: 1
     TILE_K4: 1
     TILE_N4: 1
@@ -31,6 +32,8 @@ sdpa_compute_attn_weights_tiled:
           suffix: half
   shader_variants:
     - NAME: sdpa_compute_attn_weights_tiled
+    - NAME: sdpa_compute_attn_weights_tiled_mask
+      HAS_MASK: true
     - NAME: fused_sdpa_qk_tiled
       MODE: fused
       OUT_DTYPE: float

backends/vulkan/runtime/graph/ops/glsl/sdpa_fp_attn_weight_tile_store.glslh

@@ -151,6 +151,42 @@ void apply_bias(
 }
 #endif // HAS_BIAS
 
+#ifdef HAS_MASK
+/*
+ * Adds the LLM SDPA additive attn_mask, a 2D [S, C] tensor shared by every
+ * head, to the tile. Requires t_mask and mask_sizes. The mask is read element
+ * by element: a buffer-backed mask is not padded along its width, so its rows
+ * cannot be addressed in texels unless C is a multiple of 4.
+ */
+void apply_attn_mask(
+    inout FPOutTile tile,
+    const int c_idx_start,
+    const int s_idx_start) {
+  const int mask_C = mask_sizes.x;
+  const int mask_S = mask_sizes.y;
+  [[unroll]] for (int s = 0; s < TILE_M; ++s) {
+    const int s_idx = s_idx_start + s;
+    if (s_idx >= mask_S) {
+      continue;
+    }
+    [[unroll]] for (int c4 = 0; c4 < TILE_N4; ++c4) {
+      [[unroll]] for (int c4i = 0; c4i < 4; ++c4i) {
+        const int c_idx = c_idx_start + mul_4(c4) + c4i;
+        if (c_idx < mask_C) {
+#ifdef INPUT_BUFFER
+          const float mask_val = float(t_mask[s_idx * mask_C + c_idx]);
+#else
+          const float mask_val = float(texelFetch(
+              t_mask, ivec3(div_4(c_idx), s_idx, 0), 0)[mod_4(c_idx)]);
+#endif
+          tile.data[s][c4][c4i] += T(mask_val);
+        }
+      }
+    }
+  }
+}
+#endif // HAS_MASK
+
 void set_out_tile_to_vec(out FPOutTile tile, const LINEAR_FP_OUTPUT_TILE_VEC4_T vec) {
   [[unroll]] for (int s = 0; s < TILE_M; ++s) {
     [[unroll]] for (int c4 = 0; c4 < TILE_N4; ++c4) { tile.data[s][c4] = vec; }

backends/vulkan/runtime/graph/ops/impl/SDPA.cpp

@@ -243,9 +243,14 @@ vkapi::ShaderInfo pick_sdpa_qk_shader(
     const ValueRef q_projected = args.at(1).refs.at(0);
     const ValueRef k_cache = args.at(1).refs.at(1);
     const bool is_gemv = is_single_token(graph, q_projected);
+    // An additive attn_mask is bound as a third input: q, k, attn_mask.
+    const bool has_mask = args.at(1).refs.size() >= 3;
 
     std::string shader_name = "sdpa_compute_attn_weights";
     shader_name += is_gemv ? "_coop" : "_tiled";
+    if (has_mask) {
+      shader_name += "_mask";
+    }
     add_storage_type_suffix(shader_name, graph->storage_type_of(q_projected));
     add_storage_type_suffix(shader_name, graph->storage_type_of(k_cache));
     add_dtype_suffix(shader_name, graph->dtype_of(q_projected));
@@ -481,11 +486,12 @@ void add_sdpa_kv_cache_update_node(
 }
 
 // Unified QK node (attn_weights = scale * Q @ K^T [+ bias]).
-// LLM: pass input_pos_symint (real symint), attn_mask = kDummyValueRef.
+// LLM: pass input_pos_symint (real symint), and attn_mask = a 2D [S, C]
+//      additive mask that replaces the causal mask, or kDummyValueRef.
 // FUSED: pass input_pos_symint = kDummyValueRef, attn_mask = valid ref or
 //        kDummyValueRef to indicate no bias. scale_val is always passed as
-//        a spec const; the LLM path computes it per head_dim and FUSED may
-//        inherit from the caller-supplied scale.
+//        a spec const; both modes default it to 1/sqrt(head_dim) and use the
+//        caller-supplied scale when there is one.
 void add_sdpa_compute_attn_weights_node(
     ComputeGraph& graph,
     const ValueRef q,
@@ -503,7 +509,8 @@ void add_sdpa_compute_attn_weights_node(
 
   if (mode == SDPAMode::LLM) {
     param_ubos.append(graph.get_or_create_int_param_buffer(input_pos_symint));
-  } else if (is_valid(attn_mask)) {
+  }
+  if (is_valid(attn_mask)) {
     param_ubos.append(graph.sizes_ubo(attn_mask));
     read_inputs.push_back(attn_mask);
   }
@@ -640,6 +647,33 @@ void add_sdpa_compute_out_node(
       resize_sdpa_out_node));
 }
 
+// For a masked LLM SDPA, sets out_symint = C - S (the number of keys in k
+// minus the number of queries), which the SDPA shaders read in place of
+// input_pos so that they attend to every key.
+void resize_attn_mask_input_pos_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)args;
+  const ValueRef q = resize_args.at(0);
+  const ValueRef k = resize_args.at(1);
+  const ValueRef out_symint = resize_args.at(2);
+  graph->set_symint(
+      out_symint,
+      graph->size_at<int32_t>(-3, k) - graph->size_at<int32_t>(-3, q));
+}
+
+ValueRef add_attn_mask_input_pos_node(
+    ComputeGraph& graph,
+    const ValueRef q,
+    const ValueRef k) {
+  const ValueRef out_symint = graph.add_symint(
+      graph.size_at<int32_t>(-3, k) - graph.size_at<int32_t>(-3, q));
+  graph.execute_nodes().emplace_back(
+      new ExecuteNode(resize_attn_mask_input_pos_node, {q, k, out_symint}));
+  return out_symint;
+}
+
 //
 // High level operator impl
 //
@@ -696,11 +730,30 @@ void sdpa_impl(ComputeGraph& graph, const std::vector<ValueRef>& args) {
   VK_CHECK_COND(
       graph.val_is_none(dropout_p) ||
       graph.extract_scalar<double>(dropout_p) == 0);
-  VK_CHECK_COND(graph.val_is_none(scale));
-  // is_causal is assumed to be true in the current implementation.
-  VK_CHECK_COND(
-      graph.val_is_none(is_causal) || graph.extract_scalar<bool>(is_causal));
-  VK_CHECK_COND(graph.val_is_none(attn_mask));
+
+  // As in the CPU op, an attn_mask replaces the causal mask: it is a 2D
+  // [S, C] additive mask over every key in k_cache, and cannot be combined
+  // with is_causal. Without one, the attention is assumed to be causal.
+  const bool has_mask = graph.val_is_not_none(attn_mask);
+  if (has_mask) {
+    VK_CHECK_COND(
+        graph.val_is_none(is_causal) || !graph.extract_scalar<bool>(is_causal),
+        "attn_mask and is_causal cannot both be set");
+    VK_CHECK_COND(graph.dim_of(attn_mask) == 2);
+    VK_CHECK_COND(
+        graph.size_at<int64_t>(-2, attn_mask) ==
+        graph.size_at<int64_t>(-3, q_projected));
+    VK_CHECK_COND(
+        graph.size_at<int64_t>(-1, attn_mask) ==
+        graph.size_at<int64_t>(-3, k_cache));
+    VK_CHECK_COND(graph.dtype_of(attn_mask) == graph.dtype_of(q_projected));
+    VK_CHECK_COND(
+        graph.storage_type_of(attn_mask) == graph.storage_type_of(q_projected));
+    VK_CHECK_COND(graph.packed_dim_of(attn_mask) == WHCN::kWidthDim);
+  } else {
+    VK_CHECK_COND(
+        graph.val_is_none(is_causal) || graph.extract_scalar<bool>(is_causal));
+  }
 
   const int64_t num_q_heads = graph.size_at<int64_t>(-2, q_projected);
   int64_t max_seq_len = graph.size_at<int64_t>(-3, q_projected);
@@ -763,14 +816,23 @@ void sdpa_impl(ComputeGraph& graph, const std::vector<ValueRef>& args) {
       utils::kWidthPacked);
 
   const int32_t head_dim_size = graph.size_at<int32_t>(-1, q_projected);
-  const float scale_val = 1.0f / std::sqrt(static_cast<float>(head_dim_size));
+  const float scale_val = graph.val_is_none(scale)
+      ? 1.0f / std::sqrt(static_cast<float>(head_dim_size))
+      : graph.extract_scalar<float>(scale);
+
+  // The SDPA shaders attend to the first input_pos + S keys. A masked
+  // attention attends to all C keys in k_cache, so give them C - S in place of
+  // input_pos.
+  const ValueRef attn_input_pos = has_mask
+      ? add_attn_mask_input_pos_node(graph, q_projected, k_cache)
+      : input_pos_symint;
 
   add_sdpa_compute_attn_weights_node(
       graph,
       q_projected,
       k_cache,
-      input_pos_symint,
-      /*attn_mask=*/kDummyValueRef,
+      attn_input_pos,
+      has_mask ? attn_mask : kDummyValueRef,
       scale_val,
       attn_weights,
       SDPAMode::LLM);
@@ -780,7 +842,7 @@ void sdpa_impl(ComputeGraph& graph, const std::vector<ValueRef>& args) {
       attn_weights,
       q_projected,
       k_cache,
-      input_pos_symint,
+      attn_input_pos,
       attn_weights_softmax,
       SDPAMode::LLM);
 
@@ -790,7 +852,7 @@ void sdpa_impl(ComputeGraph& graph, const std::vector<ValueRef>& args) {
       v_cache,
       q_projected,
       /*k=*/kDummyValueRef,
-      input_pos_symint,
+      attn_input_pos,
       out,
       SDPAMode::LLM);
 }