[Feature][Relax] Support shared-KV attention with configurable sliding windows

python/tvm/relax/frontend/nn/llm/_kernel_common.py

@@ -135,6 +135,15 @@ def _causal_mask(causal, row, col, kv_len, qo_len):
     )
 
 
+def _causal_or_sliding_cross_mask(causal, row, col, kv_len, qo_len, sliding_window_size):
+    visible_past = T.max(sliding_window_size - row - 1, 0)
+    return T.if_then_else(
+        tirx.And(causal > 0, sliding_window_size > 0),
+        tirx.And(col < kv_len, col >= T.max(kv_len - visible_past, 0)),
+        _causal_mask(causal, row, col, kv_len, qo_len),
+    )
+
+
 def _declare_length_info(var_length_info, batch_size, sliding_window, elem_offset):
     return (
         T.match_buffer(var_length_info, (3, batch_size), "int32", elem_offset=elem_offset)
@@ -248,7 +257,7 @@ def softmax_update_causal(
         S_smem: T.Buffer, m_smem: T.Buffer, d_smem: T.Buffer, m_prev_smem: T.Buffer,
         m_new: T.Buffer, m_prev: T.Buffer, d_new: T.Buffer,
         ty: T.int32, tx: T.int32, LH_start: T.int32, L_kv_start: T.int32,
-        causal: T.int32, kv_len: T.int32, qo_len: T.int32,
+        causal: T.int32, kv_len: T.int32, qo_len: T.int32, sliding_window_size: T.int32,
     ):
         # Phase 1: compute m_new = max(masked S over kv tile), d_new = d_prev * exp2(m_prev - m_new)
         for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
@@ -259,7 +268,7 @@ def softmax_update_causal(
                     m_new[i] = m_smem[row]
                     row_: T.let[T.int32] = (LH_start + row) // group_size
                     for j in T.serial(tile_z):
-                        if _causal_mask(causal, row=row_, col=L_kv_start + j, kv_len=kv_len, qo_len=qo_len):
+                        if _causal_or_sliding_cross_mask(causal, row=row_, col=L_kv_start + j, kv_len=kv_len, qo_len=qo_len, sliding_window_size=sliding_window_size):
                             m_new[i] = T.max(m_new[i], S_smem[row, j])
                     d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
         # Phase 2: exp-and-scale S_smem; masked-out entries use -inf
@@ -270,7 +279,7 @@ def softmax_update_causal(
                     # predicate sits inside loop so sync stays outside conditional branches
                     if row < tile_x:
                         row_: T.let[T.int32] = (LH_start + row) // group_size
-                        if _causal_mask(causal, row=row_, col=L_kv_start + j, kv_len=kv_len, qo_len=qo_len):
+                        if _causal_or_sliding_cross_mask(causal, row=row_, col=L_kv_start + j, kv_len=kv_len, qo_len=qo_len, sliding_window_size=sliding_window_size):
                             S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
                         else:
                             S_smem[row, j] = T.exp2(-5e4 - m_new[i])

python/tvm/relax/frontend/nn/llm/_prefill_kernels.py

@@ -40,6 +40,7 @@
     _alloc_softmax_state_buffers,
     _alloc_tile_walk_state,
     _causal_mask,
+    _causal_or_sliding_cross_mask,
     _declare_length_info,
     _get_kv_chunk_len,
     _get_prefill_kernel_config,
@@ -51,7 +52,14 @@
 
 
 def _attention_prefill_cpu(
-    h_kv, h_q, d, dtype, sliding_window: bool, rope_scaling: dict[str, Any], page_size: int = 16
+    h_kv,
+    h_q,
+    d,
+    dtype,
+    sliding_window: bool,
+    rope_scaling: dict[str, Any],
+    page_size: int = 16,
+    sliding_window_size: int = 1024,
 ):
     global_symbol = "batch_prefill_paged_kv_cpu"
     if sliding_window:
@@ -173,11 +181,14 @@ def batch_prefill_paged_kv_cpu(
                                 S_val[0] *= sm_scale * math.log2(math.exp(1))
 
                                 # update m_val, d_val , O_local
-                                if _causal_mask(causal,
+                                if _causal_or_sliding_cross_mask(
+                                    causal,
                                     row=q_idx,
                                     col=row_idx,
                                     kv_len=kv_chunk_len[0],
-                                    qo_len=q_indptr[b_idx + 1] - q_indptr[b_idx]):
+                                    qo_len=q_indptr[b_idx + 1] - q_indptr[b_idx],
+                                    sliding_window_size=(sliding_window_size if sliding_window else 0),
+                                ):
                                     new_m[0] = T.max(m_val[0], S_val[0])
                                 else:
                                     S_val[0] = -5e4
@@ -203,7 +214,17 @@ def batch_prefill_paged_kv_cpu(
     return batch_prefill_paged_kv_cpu
 
 
-def _attention_prefill(h_kv, h_q, d, dtype, sliding_window: bool, rope_scaling: dict[str, Any], target: Target, page_size: int = 16):
+def _attention_prefill(
+    h_kv,
+    h_q,
+    d,
+    dtype,
+    sliding_window: bool,
+    rope_scaling: dict[str, Any],
+    target: Target,
+    page_size: int = 16,
+    sliding_window_size: int = 1024,
+):
     NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z = _get_prefill_kernel_config(h_kv, h_q, d, dtype, target)
 
     global_symbol = "batch_prefill_paged_kv"
@@ -355,7 +376,7 @@ def batch_prefill_paged_kv(
                                         T.tvm_storage_sync("shared")
 
                                         compute_s_gemm(Q_smem, K_smem, S_local, S_smem, sm_scale)
-                                        softmax_update_causal(S_smem, m_smem, d_smem, m_prev_smem, m_new, m_prev, d_new, ty, tx, LH_start, L_kv_start, causal, kv_chunk_len[0], q_indptr[b_idx + 1] - q_indptr[b_idx])
+                                        softmax_update_causal(S_smem, m_smem, d_smem, m_prev_smem, m_new, m_prev, d_new, ty, tx, LH_start, L_kv_start, causal, kv_chunk_len[0], q_indptr[b_idx + 1] - q_indptr[b_idx], sliding_window_size if sliding_window else 0)
                                         compute_o_gemm(S_smem, V_smem, O_local, m_prev_smem, m_smem)
 
                                     paged_store_output_lse(output, lse, O_local, m_smem, d_smem, q_indptr, b_idx, by, LH_start)
@@ -459,7 +480,7 @@ def batch_sequence_prefill_kv(  # pylint: disable=too-many-branches
                                 T.tvm_storage_sync("shared")
 
                                 compute_s_gemm(Q_smem, K_smem, S_local, S_smem, sm_scale)
-                                softmax_update_causal(S_smem, m_smem, d_smem, m_prev_smem, m_new, m_prev, d_new, ty, tx, LH_start, L_kv_start, causal, kv_len, qo_len)
+                                softmax_update_causal(S_smem, m_smem, d_smem, m_prev_smem, m_new, m_prev, d_new, ty, tx, LH_start, L_kv_start, causal, kv_len, qo_len, 0)
                                 compute_o_gemm(S_smem, V_smem, O_local, m_prev_smem, m_smem)
 
                             # Store O from smem to gmem
@@ -889,7 +910,7 @@ def batch_prefill_ragged_kv(  # pylint: disable=too-many-branches
                                         T.tvm_storage_sync("shared")
 
                                         compute_s_gemm(Q_smem, K_smem, S_local, S_smem, sm_scale)
-                                        softmax_update_causal(S_smem, m_smem, d_smem, m_prev_smem, m_new, m_prev, d_new, ty, tx, LH_start, L_kv_start, causal, kv_chunk_len[0], q_indptr[b_idx + 1] - q_indptr[b_idx])
+                                        softmax_update_causal(S_smem, m_smem, d_smem, m_prev_smem, m_new, m_prev, d_new, ty, tx, LH_start, L_kv_start, causal, kv_chunk_len[0], q_indptr[b_idx + 1] - q_indptr[b_idx], 0)
                                         compute_o_gemm(S_smem, V_smem, O_local, m_prev_smem, m_smem)
 
                                     paged_store_output_lse(output, lse, O_local, m_smem, d_smem, q_indptr, b_idx, by, LH_start)
@@ -1030,6 +1051,7 @@ def batch_prefill_paged_kv_mla(
                                         m_new, m_prev, d_new,
                                         ty, tx, LH_start, L_kv_start,
                                         causal, kv_chunk_len[0], q_indptr[b_idx + 1] - q_indptr[b_idx],
+                                        0,
                                     )
 
                                     compute_o_gemm(S_smem, KV_smem, O_local, m_prev_smem, m_smem)

python/tvm/relax/frontend/nn/llm/kv_cache.py

@@ -230,6 +230,46 @@ def cross_attention(
         lse = Tensor(_expr=bb.emit(rx.TupleGetItem(attn_results, 1))).reshape(b, s, h_qo)
         return o, lse
 
+    def attention_with_shared_kv(
+        self,
+        source_layer_id: int,
+        q: Tensor,
+        current_k: Tensor,
+        current_v: Tensor,
+        sm_scale: float,
+    ) -> Tensor:
+        """Compute attention using K/V shared by another logical layer.
+
+        This operation does not append to the KV cache. The cache-owning source layer must run
+        before this operation in the active forward pass. ``current_k`` and ``current_v`` are that
+        source layer's K/V for the active chunk. During prefill they are combined with the source
+        layer's paged past K/V; during decode the source layer has already appended them and the
+        operation reads the cache directly.
+        """
+        # pylint: disable=protected-access
+        b, s, h_qo, d_qk = q._expr.ty.shape
+        _, _, h_kv, _ = current_k._expr.ty.shape
+        _, _, _, d_v = current_v._expr.ty.shape
+        q = q.reshape(b * s, h_qo, d_qk)
+        current_k = current_k.reshape(b * s, h_kv, d_qk)
+        current_v = current_v.reshape(b * s, h_kv, d_v)
+        return Tensor(
+            _expr=rx.BlockBuilder.current().emit(
+                rx.call_dps_packed(
+                    "vm.builtin.attention_kv_cache_attention_with_shared_kv",
+                    [
+                        self._expr,
+                        rx.prim_value(source_layer_id),  # type: ignore[arg-type]
+                        rx.prim_value(sm_scale),
+                        q._expr,
+                        current_k._expr,
+                        current_v._expr,
+                    ],
+                    out_ty=rx.TensorType((b * s, h_qo, d_v), q.dtype),
+                )
+            )
+        ).reshape(b, s, h_qo, d_v)
+
     def append_mla_kv(self, layer_id: int, kv: Tensor) -> "PagedKVCache":
         """Fine-grained API that appends the MLA K/V data to KV cache."""
         # pylint: disable=protected-access
@@ -355,6 +395,8 @@ def __init__(  # pylint: disable=too-many-locals
         dtype: str,
         target: Target,
         name: str = "paged_kv_cache",
+        *,
+        layer_sliding_window_size: int = 1024,
     ) -> None:
         """Create a paged KV cache object with FlashInfer kernels.
 
@@ -399,6 +441,8 @@ def __init__(  # pylint: disable=too-many-locals
             The number of dimensions in the embedding that RoPE is applied to.
         enable_disaggregation : bool
             Whether to enable disaggregation in the KV cache.
+        layer_sliding_window_size : int
+            Window size used by layers whose attention kind is ``mha_sliding``.
         """
         assert rope_mode != RopeMode.INLINE, "FlashInfer RoPE does not support inline mode."
         rope_scaling = _prepare_yarn_rope_scaling(rope_scaling, rope_theta)
@@ -447,7 +491,7 @@ def __init__(  # pylint: disable=too-many-locals
             [
                 rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_prefill_paged_run"), rx.ExternFunc("batch_prefill_plan")]),
                 rx.Tuple([rx.StringImm("flashinfer"), rx.ExternFunc("batch_decode_run"), rx.ExternFunc("batch_decode_plan")]),
-                rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_prefill_sliding_window")]),
+                rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target, sliding_window_size=layer_sliding_window_size), "tir_attention_prefill_sliding_window")]),
                 rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
                 rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
                 rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),
@@ -476,6 +520,7 @@ def __init__(  # pylint: disable=too-many-locals
                     prefill_chunk_size,
                     page_size,
                     support_sliding_window,
+                    layer_sliding_window_size,
                 ]
             ),
             layer_partition,
@@ -540,6 +585,8 @@ def __init__(  # pylint: disable=too-many-locals
         dtype: str,
         target: Target,
         name: str = "paged_kv_cache",
+        *,
+        layer_sliding_window_size: int = 1024,
     ) -> None:
         """Create a paged KV cache object with TIR kernels.
 
@@ -586,6 +633,8 @@ def __init__(  # pylint: disable=too-many-locals
             Whether to enable disaggregation in the KV cache.
         target : Target
             The target to build the model to.
+        layer_sliding_window_size : int
+            Window size used by layers whose attention kind is ``mha_sliding``.
         """
         rope_scaling = _prepare_yarn_rope_scaling(rope_scaling, rope_theta)
         attn_kind_single = attn_kind[0] if isinstance(attn_kind, list) else attn_kind
@@ -604,6 +653,7 @@ def __init__(  # pylint: disable=too-many-locals
                     prefill_chunk_size,
                     page_size,
                     support_sliding_window,
+                    layer_sliding_window_size,
                 ]
             ),
             layer_partition,
@@ -630,7 +680,7 @@ def __init__(  # pylint: disable=too-many-locals
                     rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_ragged_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, v_head_dim, dtype, rope_scaling), "tir_attention_prefill_ragged_cpu")]),
                     rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_prefill_cpu")]),
                     rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling), "tir_attention_decode_cpu")]),
-                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling), "tir_attention_prefill_cpu_sliding_window")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, sliding_window_size=layer_sliding_window_size), "tir_attention_prefill_cpu_sliding_window")]),
                     rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling), "tir_attention_decode_cpu_sliding_window")]),
                     rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_cpu")]),
                     rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache_cpu(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache_cpu")]),
@@ -650,7 +700,7 @@ def __init__(  # pylint: disable=too-many-locals
                 [
                     rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_prefill")]),
                     rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, False, rope_scaling, target), "tir_attention_decode")]),
-                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_prefill_sliding_window")]),
+                    rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_prefill(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target, sliding_window_size=layer_sliding_window_size), "tir_attention_prefill_sliding_window")]),
                     rx.Tuple([rx.StringImm("tirx"), bb.add_func(_attention_decode(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, True, rope_scaling, target), "tir_attention_decode_sliding_window")]),
                     rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn_with_paged_kv_cache(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask_with_paged_kv_cache")]),
                     rx.Tuple([rx.StringImm("tirx"), bb.add_func(tree_attn(num_key_value_heads, num_attention_heads, qk_head_dim, dtype, rope_scaling, target), "tir_attention_prefill_with_tree_mask")]),

src/runtime/vm/kv_state.cc

@@ -94,6 +94,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
              kv_cache->CrossAttention(layer_id, std::move(q_data), std::move(o_data),
                                       std::move(lse_data), sm_scale);
            })
+      .def("vm.builtin.attention_kv_cache_attention_with_shared_kv",
+           [](AttentionKVCache kv_cache, int64_t source_layer_id, double sm_scale, Tensor q_data,
+              Tensor current_k_data, Tensor current_v_data, Tensor o_data) {
+             kv_cache->AttentionWithSharedKV(source_layer_id, std::move(q_data),
+                                             std::move(current_k_data), std::move(current_v_data),
+                                             std::move(o_data), sm_scale);
+           })
       .def("vm.builtin.attention_kv_cache_append_mla_kv",
            [](AttentionKVCache kv_cache, int64_t layer_id, Tensor kv_data) {
              kv_cache->AppendMLAKV(layer_id, std::move(kv_data));

src/runtime/vm/kv_state.h

@@ -206,6 +206,21 @@ class AttentionKVCacheObj : public KVStateObj {
   virtual void CrossAttention(int64_t layer_id, Tensor q_data, Tensor o_data, Tensor lse_data,
                               double sm_scale) = 0;
 
+  /*!
+   * \brief Compute attention with K/V shared by another logical layer.
+   *
+   * This operation does not append K/V. The cache-owning source layer must run before this
+   * operation in the active forward pass.
+   * \param source_layer_id The physical cache layer containing past K/V data.
+   * \param q_data The logical layer's input Q data.
+   * \param current_k_data The source layer's K data for the active chunk.
+   * \param current_v_data The source layer's V data for the active chunk.
+   * \param o_data The output O data.
+   * \param sm_scale The additional attention scaling factor.
+   */
+  virtual void AttentionWithSharedKV(int64_t source_layer_id, Tensor q_data, Tensor current_k_data,
+                                     Tensor current_v_data, Tensor o_data, double sm_scale) = 0;
+
   /*!
    * \brief Fine-grained API that appends the MLA K/V data to KV cache.
    * \param layer_id The model layer where the attention compute happens.