This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 27c2e019d0 [Feature][Relax] Support shared-KV attention with
configurable sliding windows (#20121)
27c2e019d0 is described below
commit 27c2e019d0ce6182158020c7534dda4a3ce981ae
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Wed Aug 12 14:19:27 2026 -0700
[Feature][Relax] Support shared-KV attention with configurable sliding
windows (#20121)
Extend `PagedKVCache` for models whose logical attention layers reuse
K/V cached by another physical layer.
- Adds non-mutating `attention_with_shared_kv` support.
- Makes per-layer sliding-window size configurable.
- Corrects sliding-window masking during chunked prefill.
This is intended to enable downstream Gemma 4 support in MLC-LLM and
WebLLM.
---
python/tvm/relax/frontend/nn/llm/_kernel_common.py | 15 +-
.../tvm/relax/frontend/nn/llm/_prefill_kernels.py | 36 ++++-
python/tvm/relax/frontend/nn/llm/kv_cache.py | 56 ++++++-
src/runtime/vm/kv_state.cc | 7 +
src/runtime/vm/kv_state.h | 15 ++
src/runtime/vm/paged_kv_cache.cc | 131 ++++++++++++-----
...runtime_builtin_paged_attention_kv_cache_cpu.py | 161 +++++++++++++++++++--
...runtime_builtin_paged_attention_kv_cache_tir.py | 65 +++++++++
8 files changed, 422 insertions(+), 64 deletions(-)
diff --git a/python/tvm/relax/frontend/nn/llm/_kernel_common.py
b/python/tvm/relax/frontend/nn/llm/_kernel_common.py
index 6d7450e4fa..2de72360ee 100644
--- a/python/tvm/relax/frontend/nn/llm/_kernel_common.py
+++ b/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 _make_prefill_macros(tile_x, tile_y, tile_z, tile_o,
bdx, num_warps, group_s
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 _make_prefill_macros(tile_x, tile_y, tile_z, tile_o,
bdx, num_warps, group_s
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 _make_prefill_macros(tile_x, tile_y, tile_z, tile_o,
bdx, num_warps, group_s
# 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])
diff --git a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
index 54e7424e71..9eeb18f8cf 100644
--- a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
+++ b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
@@ -40,6 +40,7 @@ from ._kernel_common import (
_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 @@ from ._kernel_common import (
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 _attention_prefill_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 _attention_prefill_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 _attention_prefill(h_kv, h_q, d, dtype, sliding_window:
bool, rope_scaling:
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 _attention_sequence_prefill(h_kv, h_q, d, dtype,
target: Target, causal=0, s
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 _attention_prefill_ragged(h_kv, h_q, d_qk, d_v, dtype,
rope_scaling: dict[st
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 _attention_prefill_mla(h_q, d_latent, d_rope, dtype,
sliding_window: bool, t
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)
diff --git a/python/tvm/relax/frontend/nn/llm/kv_cache.py
b/python/tvm/relax/frontend/nn/llm/kv_cache.py
index d7c35279c6..55f2491a4a 100644
--- a/python/tvm/relax/frontend/nn/llm/kv_cache.py
+++ b/python/tvm/relax/frontend/nn/llm/kv_cache.py
@@ -230,6 +230,46 @@ class PagedKVCache(Object): # pylint:
disable=too-few-public-methods
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 @@ class FlashInferPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-me
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 @@ class FlashInferPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-me
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 @@ class FlashInferPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-me
[
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 @@ class FlashInferPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-me
prefill_chunk_size,
page_size,
support_sliding_window,
+ layer_sliding_window_size,
]
),
layer_partition,
@@ -540,6 +585,8 @@ class TIRPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-methods
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 @@ class TIRPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-methods
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 @@ class TIRPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-methods
prefill_chunk_size,
page_size,
support_sliding_window,
+ layer_sliding_window_size,
]
),
layer_partition,
@@ -630,7 +680,7 @@ class TIRPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-methods
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 @@ class TIRPagedKVCache(PagedKVCache): # pylint:
disable=too-few-public-methods
[
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")]),
diff --git a/src/runtime/vm/kv_state.cc b/src/runtime/vm/kv_state.cc
index 05f951dc74..2d93563497 100644
--- a/src/runtime/vm/kv_state.cc
+++ b/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));
diff --git a/src/runtime/vm/kv_state.h b/src/runtime/vm/kv_state.h
index 198bd18d97..4e1694c899 100644
--- a/src/runtime/vm/kv_state.h
+++ b/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.
diff --git a/src/runtime/vm/paged_kv_cache.cc b/src/runtime/vm/paged_kv_cache.cc
index efd099e826..cd0722f865 100644
--- a/src/runtime/vm/paged_kv_cache.cc
+++ b/src/runtime/vm/paged_kv_cache.cc
@@ -103,6 +103,8 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj
{
const bool support_sliding_window_;
/*! \brief A boolean flag indicating if the KV cache has per layer sliding
window. */
const bool support_layer_sliding_window_;
+ /*! \brief Window size used by layers with per-layer sliding attention. */
+ const int64_t layer_sliding_window_size_;
/*! \brief The attention kinds for each layer. */
const std::vector<AttnKind> attn_kinds_;
@@ -284,6 +286,22 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
/*! \brief The device stream for KV transfer */
TVMStreamHandle kv_transfer_stream_ = nullptr;
+ int32_t GetLayerSlidingWindowOffset(int64_t seq_length) const {
+ if (seq_length <= layer_sliding_window_size_) {
+ return 0;
+ }
+ return static_cast<int32_t>((seq_length - layer_sliding_window_size_) %
page_size_);
+ }
+
+ int32_t GetLayerSlidingWindowNumPages(int64_t seq_length) const {
+ if (seq_length == 0) {
+ return 0;
+ }
+ int64_t window_length = std::min(seq_length, layer_sliding_window_size_);
+ return static_cast<int32_t>(
+ (GetLayerSlidingWindowOffset(seq_length) + window_length + page_size_
- 1) / page_size_);
+ }
+
public:
/*! \brief Constructor. Take the cache configuration and initialize the
Tensors. */
explicit PagedAttentionKVCacheObj(
@@ -291,9 +309,9 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj
{
int64_t layer_id_end_offset, int64_t num_qo_heads, int64_t num_kv_heads,
int64_t qk_head_dim,
int64_t v_head_dim, std::vector<AttnKind> attn_kinds, int64_t
reserved_num_seqs,
int64_t num_total_pages, int64_t prefill_chunk_size, bool
support_sliding_window,
- RoPEMode rope_mode, double rotary_scale, double rotary_theta,
- ffi::Optional<Tensor> rope_ext_factors, bool enable_kv_transfer,
DLDataType dtype,
- Device device, ffi::Optional<ffi::Function> f_transpose_append_mha,
+ int64_t layer_sliding_window_size, RoPEMode rope_mode, double
rotary_scale,
+ double rotary_theta, ffi::Optional<Tensor> rope_ext_factors, bool
enable_kv_transfer,
+ DLDataType dtype, Device device, ffi::Optional<ffi::Function>
f_transpose_append_mha,
ffi::Optional<ffi::Function> f_transpose_append_mla, ffi::Function
f_compact_copy,
std::unique_ptr<RaggedPrefillFunc> f_attention_prefill_ragged,
std::unique_ptr<PagedPrefillFunc> f_attention_prefill,
@@ -320,6 +338,7 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj
{
: support_sliding_window),
support_layer_sliding_window_(std::find(attn_kinds.begin(),
attn_kinds.end(),
AttnKind::kMHASliding) !=
attn_kinds.end()),
+ layer_sliding_window_size_(layer_sliding_window_size),
attn_kinds_(std::move(attn_kinds)),
rope_mode_(support_sliding_window && rope_mode != RoPEMode::kNone ?
RoPEMode::kInline
:
rope_mode),
@@ -345,6 +364,8 @@ class PagedAttentionKVCacheObj : public AttentionKVCacheObj
{
f_copy_single_page_(std::move(f_copy_single_page)),
f_debug_get_kv_(std::move(f_debug_get_kv)),
device_(device) {
+ TVM_FFI_ICHECK_GT(layer_sliding_window_size_, 0)
+ << "Per-layer sliding window size must be positive.";
// Note: For MLA, sliding window and disaggregation are disabled for now.
if (std::find(attn_kinds_.begin(), attn_kinds_.end(), AttnKind::kMLA) !=
attn_kinds_.end()) {
TVM_FFI_ICHECK(!support_sliding_window_) << "Sliding window not
supported yet for MLA";
@@ -1024,12 +1045,12 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
}
// For sliding window, the first page and last page will both be
partially used
- page_indptr_sliding_window_h.push_back(
- page_indptr_sliding_window_h.back() +
+ int32_t num_layer_sliding_pages =
std::min(static_cast<int32_t>(block.page_ids.size()),
- static_cast<int32_t>(1024 / page_size_ +
- (block.seq_length % page_size_ ?
1 : 0))));
- for (int i = page_indices_h.size() -
page_indptr_sliding_window_h.back();
+ GetLayerSlidingWindowNumPages(block.seq_length));
+
page_indptr_sliding_window_h.push_back(page_indptr_sliding_window_h.back() +
+ num_layer_sliding_pages);
+ for (int i = page_indices_h.size() - num_layer_sliding_pages;
i < static_cast<int32_t>(page_indices_h.size()); i++) {
page_indices_sliding_window_h.push_back(page_indices_h[i]);
}
@@ -1042,11 +1063,7 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
page_size_ +
1);
if (support_layer_sliding_window_) {
- if (block.seq_length < 1024) {
- sliding_window_offset_h.push_back(0);
- } else {
- sliding_window_offset_h.push_back(block.seq_length %
page_size_);
- }
+
sliding_window_offset_h.push_back(GetLayerSlidingWindowOffset(block.seq_length));
} else {
sliding_window_offset_h.push_back(block.sliding_window_offset);
}
@@ -1055,8 +1072,8 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
// If sliding window, we need to calculate the positional offset
if (support_layer_sliding_window_) {
- k_rope_pos_offset_sliding_window_h.push_back(
- std::max(0, block.start_pos + block.seq_length - 1024));
+ k_rope_pos_offset_sliding_window_h.push_back(std::max<int64_t>(
+ 0, block.start_pos + block.seq_length -
layer_sliding_window_size_));
}
} else {
// Blocks at maximum depth
@@ -1078,12 +1095,11 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
last_block_id = id;
}
page_indptr_h.push_back(page_indptr_h.back() + num_pages);
- page_indptr_sliding_window_h.push_back(
- page_indptr_sliding_window_h.back() +
- std::min(static_cast<int32_t>(block.page_ids.size()),
- static_cast<int32_t>(1024 / page_size_ +
- (block.seq_length % page_size_ ?
1 : 0))));
- for (int i = page_indices_h.size() -
page_indptr_sliding_window_h.back();
+ int32_t num_layer_sliding_pages =
+ std::min(num_pages,
GetLayerSlidingWindowNumPages(total_seq_length));
+
page_indptr_sliding_window_h.push_back(page_indptr_sliding_window_h.back() +
+ num_layer_sliding_pages);
+ for (int i = page_indices_h.size() - num_layer_sliding_pages;
i < static_cast<int32_t>(page_indices_h.size()); i++) {
page_indices_sliding_window_h.push_back(page_indices_h[i]);
}
@@ -1095,19 +1111,15 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
page_size_ +
1);
if (support_layer_sliding_window_) {
- if (last_block.seq_length < 1024) {
- sliding_window_offset_h.push_back(0);
- } else {
- sliding_window_offset_h.push_back(last_block.seq_length %
page_size_);
- }
+
sliding_window_offset_h.push_back(GetLayerSlidingWindowOffset(total_seq_length));
} else {
sliding_window_offset_h.push_back(last_block.sliding_window_offset);
}
sink_size_h.push_back(last_block.sink_length);
k_rope_pos_offset_h.push_back(block.start_pos);
if (support_layer_sliding_window_) {
- k_rope_pos_offset_sliding_window_h.push_back(
- std::max(0, block.start_pos + block.seq_length - 1024));
+ k_rope_pos_offset_sliding_window_h.push_back(std::max<int64_t>(
+ 0, block.start_pos + total_seq_length -
layer_sliding_window_size_));
}
}
}
@@ -1466,12 +1478,52 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
if (attn_kind == AttnKind::kMHA) {
MHACrossAttnInternal(local_layer_id, q_data, o_data, lse_data, sm_scale,
- /*is_first_kernel=*/true);
+ /*is_first_kernel=*/true, /*causal=*/false);
} else {
MLACrossAttnInternal(local_layer_id, q_data, o_data, lse_data, sm_scale);
}
}
+ void AttentionWithSharedKV(int64_t source_layer_id, Tensor q_data, Tensor
current_k_data,
+ Tensor current_v_data, Tensor o_data, double
sm_scale) final {
+ int64_t local_layer_id = source_layer_id - layer_id_begin_offset_;
+ TVM_FFI_ICHECK_GE(local_layer_id, 0);
+ TVM_FFI_ICHECK_LT(local_layer_id, num_layers_);
+ Tensor pages = pages_[local_layer_id];
+ TVM_FFI_ICHECK(q_data.DataType() == pages.DataType());
+ TVM_FFI_ICHECK(current_k_data.DataType() == pages.DataType());
+ TVM_FFI_ICHECK(current_v_data.DataType() == pages.DataType());
+ TVM_FFI_ICHECK(o_data.DataType() == pages.DataType());
+ TVM_FFI_ICHECK(attn_kinds_[source_layer_id] == AttnKind::kMHA ||
+ attn_kinds_[source_layer_id] == AttnKind::kMHASliding)
+ << "Querying K/V from another logical layer is only supported for MHA
caches.";
+
+ int64_t total_seq_length = 0;
+ for (int64_t seq_id = 0; seq_id < cur_batch_size_; ++seq_id) {
+ total_seq_length += cur_append_lengths_[seq_id];
+ }
+ TVM_FFI_ICHECK_EQ(q_data->ndim, 3);
+ TVM_FFI_ICHECK_EQ(current_k_data->ndim, 3);
+ TVM_FFI_ICHECK_EQ(current_v_data->ndim, 3);
+ TVM_FFI_ICHECK_EQ(o_data->ndim, 3);
+ TVM_FFI_ICHECK_EQ(q_data->shape[0], total_seq_length);
+ TVM_FFI_ICHECK_EQ(current_k_data->shape[0], total_seq_length);
+ TVM_FFI_ICHECK_EQ(current_v_data->shape[0], total_seq_length);
+ TVM_FFI_ICHECK_EQ(o_data->shape[0], total_seq_length);
+ TVM_FFI_ICHECK_EQ(q_data->shape[1], num_qo_heads_);
+ TVM_FFI_ICHECK_EQ(current_k_data->shape[1], num_kv_heads_);
+ TVM_FFI_ICHECK_EQ(current_v_data->shape[1], num_kv_heads_);
+ TVM_FFI_ICHECK_EQ(o_data->shape[1], num_qo_heads_);
+ TVM_FFI_ICHECK_EQ(q_data->shape[2], qk_head_dim_);
+ TVM_FFI_ICHECK_EQ(current_k_data->shape[2], qk_head_dim_);
+ TVM_FFI_ICHECK_EQ(current_v_data->shape[2], v_head_dim_);
+ TVM_FFI_ICHECK_EQ(o_data->shape[2], v_head_dim_);
+
+ ComputeStreamWaitForCopyStream();
+ TVM_FFI_ICHECK(!dirty_aux_data_device_);
+ AttentionInternal(source_layer_id, q_data, current_k_data, current_v_data,
o_data, sm_scale);
+ }
+
void AppendMLAKV(int64_t layer_id, Tensor kv_data) final {
// Shape and dtype check.
int64_t local_layer_id = layer_id - layer_id_begin_offset_;
@@ -2120,7 +2172,9 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
}
bool self_attn_computed = !is_first_kernel;
bool cross_attn_computed = MHACrossAttnInternal(
- local_layer_id, q_data, output, merged_attn_lse_view_, sm_scale,
is_first_kernel);
+ local_layer_id, q_data, output, merged_attn_lse_view_, sm_scale,
is_first_kernel,
+ /*causal=*/!append_before_attn_ &&
+ attn_kinds_[local_layer_id + layer_id_begin_offset_] ==
AttnKind::kMHASliding);
TVM_FFI_ICHECK(self_attn_computed || cross_attn_computed)
<< "Both self-attention and cross-attention are not computed.";
}
@@ -2160,7 +2214,7 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
/*! \brief Compute cross-attention for MHA. Return if there is effective
computation. */
bool MHACrossAttnInternal(int64_t local_layer_id, Tensor q_data, Tensor
o_data, Tensor lse_data,
- double sm_scale, bool is_first_kernel) {
+ double sm_scale, bool is_first_kernel, bool
causal) {
std::unique_ptr<PagedPrefillFunc>& f_prefill =
(!support_sliding_window_ &&
attn_kinds_[local_layer_id + layer_id_begin_offset_] !=
AttnKind::kMHASliding)
@@ -2229,8 +2283,7 @@ class PagedAttentionKVCacheObj : public
AttentionKVCacheObj {
// Use prefill kernel for depth d
TVM_FFI_ICHECK_NOTNULL(f_prefill);
f_prefill->MHA(d, q_data, qo_indptr_on_depths_view_[d],
pages_[local_layer_id], page_indptr,
- page_indices, length_info, q_rope_position_map_view_,
k_rope_pos,
- /*causal=*/false,
+ page_indices, length_info, q_rope_position_map_view_,
k_rope_pos, causal,
/*rotary_mode=*/rope_mode_, rotary_scale, rotary_theta,
sm_scale,
attn_output, attn_lse, compute_stream_);
}
@@ -2550,12 +2603,15 @@ TVM_FFI_STATIC_INIT_BLOCK() {
attn_kinds_vec.push_back(static_cast<AttnKind>(attn_kind));
}
- TVM_FFI_ICHECK_EQ(cache_config.size(), 5);
+ TVM_FFI_ICHECK(cache_config.size() == 5 || cache_config.size() == 6)
+ << "KV cache config must contain five legacy fields and an
optional per-layer "
+ "sliding window size.";
int64_t reserved_num_seqs = cache_config[0];
int64_t total_token_capacity = cache_config[1];
int64_t prefill_chunk_size = cache_config[2];
int64_t page_size = cache_config[3];
bool support_sliding_window = cache_config[4];
+ int64_t layer_sliding_window_size = cache_config.size() == 6 ?
cache_config[5] : 1024;
int64_t num_total_pages = (total_token_capacity + page_size - 1) /
page_size + 1;
if (support_sliding_window) {
// When sliding window is enabled, each sequence may use two more
pages at most.
@@ -2566,9 +2622,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::ObjectPtr<PagedAttentionKVCacheObj> n =
ffi::make_object<PagedAttentionKVCacheObj>(
page_size, num_layers, layer_id_begin_offset, layer_id_end_offset,
num_qo_heads,
num_kv_heads, qk_head_dim, v_head_dim, attn_kinds_vec,
reserved_num_seqs,
- num_total_pages, prefill_chunk_size, support_sliding_window,
RoPEMode(rope_mode),
- rotary_scale, rotary_theta, std::move(rope_ext_factors),
enable_kv_transfer, //
- init->dtype, init->device,
//
+ num_total_pages, prefill_chunk_size, support_sliding_window,
layer_sliding_window_size,
+ RoPEMode(rope_mode), rotary_scale, rotary_theta,
std::move(rope_ext_factors),
+ enable_kv_transfer, //
+ init->dtype, init->device, //
std::move(f_transpose_append_mha),
std::move(f_transpose_append_mla),
std::move(f_compact_copy), std::move(f_attention_prefill_ragged),
std::move(f_attention_prefill), std::move(f_attention_decode),
diff --git
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
index dfd0135178..ed86585e5d 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_cpu.py
@@ -71,6 +71,7 @@ fbegin_forward = None
fend_forward = None
fcommit_accepted_token_tree_nodes = None
fattention_with_fuse_qkv = None
+fattention_with_shared_kv = None
fis_empty = None
fdebug_get_kv = None
@@ -92,10 +93,10 @@ fcompact_copy = None
_COMPILED_KERNEL_CACHE = {}
-def set_global_func(head_dim, dtype):
+def set_global_func(head_dim, dtype, layer_sliding_window_size=1024):
global fclear, fadd_sequence, fremove_sequence, ffork_sequence,
fenable_sliding_window_for_seq
global fpopn, fbegin_forward, fend_forward,
fcommit_accepted_token_tree_nodes
- global fattention_with_fuse_qkv, fis_empty, fdebug_get_kv
+ global fattention_with_fuse_qkv, fattention_with_shared_kv, fis_empty,
fdebug_get_kv
global ftranspose_append, fcopy_cache, fattn_prefill, fattn_decode
global \
fattn_prefill_ragged, \
@@ -120,6 +121,9 @@ def set_global_func(head_dim, dtype):
fattention_with_fuse_qkv = tvm.get_global_func(
"vm.builtin.attention_kv_cache_attention_with_fused_qkv"
)
+ fattention_with_shared_kv = tvm.get_global_func(
+ "vm.builtin.attention_kv_cache_attention_with_shared_kv"
+ )
fis_empty = tvm.get_global_func("vm.builtin.attention_kv_cache_empty")
fdebug_get_kv =
tvm.get_global_func("vm.builtin.attention_kv_cache_debug_get_kv")
@@ -135,6 +139,7 @@ def set_global_func(head_dim, dtype):
rope_scale,
rope_theta,
json.dumps(rope_scaling, sort_keys=True),
+ layer_sliding_window_size,
)
builts = _COMPILED_KERNEL_CACHE.get(cache_key)
if builts is None:
@@ -146,7 +151,15 @@ def set_global_func(head_dim, dtype):
num_kv_heads, num_qo_heads, head_dim, dtype, False,
rope_scaling
),
_attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, dtype,
False, rope_scaling),
- _attention_prefill_cpu(num_kv_heads, num_qo_heads, head_dim,
dtype, True, rope_scaling),
+ _attention_prefill_cpu(
+ num_kv_heads,
+ num_qo_heads,
+ head_dim,
+ dtype,
+ True,
+ rope_scaling,
+ sliding_window_size=layer_sliding_window_size,
+ ),
_attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, dtype,
True, rope_scaling),
_attention_prefill_ragged_cpu(
num_kv_heads, num_qo_heads, head_dim, head_dim, dtype,
rope_scaling
@@ -187,24 +200,34 @@ def set_global_func(head_dim, dtype):
) = builts
-def create_kv_cache(head_dim, dtype, rope_mode, support_sliding_window):
+def create_kv_cache(
+ head_dim,
+ dtype,
+ rope_mode,
+ support_sliding_window,
+ layer_sliding_window_size=None,
+ attn_kinds=None,
+):
fcreate = tvm.get_global_func("vm.builtin.paged_attention_kv_cache_create")
+ cache_config = [
+ reserved_nseq,
+ maximum_total_seq_length,
+ prefill_chunk_size,
+ page_size,
+ int(support_sliding_window),
+ ]
+ if layer_sliding_window_size is not None:
+ cache_config.append(layer_sliding_window_size)
+ if attn_kinds is None:
+ attn_kinds = [int(AttnKind.MHA) for _ in range(num_layers)]
cache = fcreate(
- tvm_ffi.Shape(
- [
- reserved_nseq,
- maximum_total_seq_length,
- prefill_chunk_size,
- page_size,
- int(support_sliding_window),
- ]
- ),
+ tvm_ffi.Shape(cache_config),
tvm_ffi.Shape([0, num_layers]),
num_qo_heads,
num_kv_heads,
head_dim,
head_dim, # v_head_dim
- tvm_ffi.Shape([int(AttnKind.MHA) for _ in range(num_layers)]),
+ tvm_ffi.Shape(attn_kinds),
False, # enable_kv_transfer
rope_mode,
rope_scale,
@@ -570,6 +593,116 @@ def apply_attention(
verify_cached_kv(kv_cache, seq_ids, cached_k, cached_v)
+def _causal_attention_reference(q, k, v, past_length,
sliding_window_size=None):
+ k = np.repeat(k, num_qo_heads // num_kv_heads, axis=1)
+ v = np.repeat(v, num_qo_heads // num_kv_heads, axis=1)
+ scores = np.einsum("qhd,khd->hqk", q.astype("float32"),
k.astype("float32")) * sm_scale
+ query_positions = past_length + np.arange(q.shape[0])
+ key_positions = np.arange(k.shape[0])
+ mask = key_positions[None, :] <= query_positions[:, None]
+ if sliding_window_size is not None:
+ mask &= key_positions[None, :] > query_positions[:, None] -
sliding_window_size
+ scores = np.where(mask[None, :, :], scores, np.finfo("float32").min)
+ probabilities = scipy.special.softmax(scores, axis=-1)
+ return np.einsum("hqk,khd->qhd", probabilities, v.astype("float32"))
+
+
+def test_per_layer_sliding_window():
+ global head_dim, sm_scale, dtype
+ head_dim = 64
+ sm_scale = head_dim ** (-0.5)
+ dtype = "float32"
+ layer_window_size = 3
+ set_global_func(head_dim, dtype, layer_window_size)
+ kv_cache = create_kv_cache(
+ head_dim,
+ dtype,
+ RopeMode.NONE,
+ False,
+ layer_sliding_window_size=layer_window_size,
+ attn_kinds=[int(AttnKind.MHA_SLIDING)] + [int(AttnKind.MHA)] *
(num_layers - 1),
+ )
+ fadd_sequence(kv_cache, 0)
+
+ rng = np.random.default_rng(0)
+ cached_k = np.empty((0, num_kv_heads, head_dim), dtype=dtype)
+ cached_v = np.empty((0, num_kv_heads, head_dim), dtype=dtype)
+ for append_length in [3, 2, 1]:
+ past_length = cached_k.shape[0]
+ q = rng.standard_normal((append_length, num_qo_heads,
head_dim)).astype(dtype)
+ current_k = rng.standard_normal((append_length, num_kv_heads,
head_dim)).astype(dtype)
+ current_v = rng.standard_normal((append_length, num_kv_heads,
head_dim)).astype(dtype)
+
+ fbegin_forward(kv_cache, Shape([0]), Shape([append_length]), None)
+ qkv = tvm.runtime.tensor(np.concatenate([q, current_k, current_v],
axis=1), device)
+ output = tvm.runtime.empty(q.shape, dtype, device=device)
+ fattention_with_fuse_qkv(kv_cache, 0, sm_scale, qkv, output)
+
+ cached_k = np.concatenate([cached_k, current_k], axis=0)
+ cached_v = np.concatenate([cached_v, current_v], axis=0)
+ expected = _causal_attention_reference(
+ q, cached_k, cached_v, past_length, layer_window_size
+ )
+ tvm.testing.assert_allclose(output.numpy(), expected, rtol=1e-3,
atol=1e-3)
+ fend_forward(kv_cache)
+
+
[email protected](
+ ("source_attn_kind", "layer_window_size"),
+ [(AttnKind.MHA, None), (AttnKind.MHA_SLIDING, 3)],
+)
+def test_attention_with_shared_kv(source_attn_kind, layer_window_size):
+ global head_dim, sm_scale, dtype
+ head_dim = 64
+ sm_scale = head_dim ** (-0.5)
+ dtype = "float32"
+ set_global_func(head_dim, dtype, layer_window_size or 1024)
+ attn_kinds = [int(source_attn_kind)] + [int(AttnKind.MHA)] * (num_layers -
1)
+ kv_cache = create_kv_cache(
+ head_dim,
+ dtype,
+ RopeMode.NONE,
+ False,
+ layer_sliding_window_size=layer_window_size or 1024,
+ attn_kinds=attn_kinds,
+ )
+ fadd_sequence(kv_cache, 0)
+
+ rng = np.random.default_rng(0)
+ cached_k = np.empty((0, num_kv_heads, head_dim), dtype=dtype)
+ cached_v = np.empty((0, num_kv_heads, head_dim), dtype=dtype)
+ for append_length in [3, 2, 1]:
+ past_length = cached_k.shape[0]
+ source_q = rng.standard_normal((append_length, num_qo_heads,
head_dim)).astype(dtype)
+ current_k = rng.standard_normal((append_length, num_kv_heads,
head_dim)).astype(dtype)
+ current_v = rng.standard_normal((append_length, num_kv_heads,
head_dim)).astype(dtype)
+ shared_q = rng.standard_normal((append_length, num_qo_heads,
head_dim)).astype(dtype)
+
+ fbegin_forward(kv_cache, Shape([0]), Shape([append_length]), None)
+ qkv = tvm.runtime.tensor(np.concatenate([source_q, current_k,
current_v], axis=1), device)
+ source_output = tvm.runtime.empty(source_q.shape, dtype, device=device)
+ fattention_with_fuse_qkv(kv_cache, 0, sm_scale, qkv, source_output)
+
+ shared_output = tvm.runtime.empty(shared_q.shape, dtype, device=device)
+ fattention_with_shared_kv(
+ kv_cache,
+ 0,
+ sm_scale,
+ tvm.runtime.tensor(shared_q, device),
+ tvm.runtime.tensor(current_k, device),
+ tvm.runtime.tensor(current_v, device),
+ shared_output,
+ )
+
+ cached_k = np.concatenate([cached_k, current_k], axis=0)
+ cached_v = np.concatenate([cached_v, current_v], axis=0)
+ expected = _causal_attention_reference(
+ shared_q, cached_k, cached_v, past_length, layer_window_size
+ )
+ tvm.testing.assert_allclose(shared_output.numpy(), expected,
rtol=1e-3, atol=1e-3)
+ fend_forward(kv_cache)
+
+
def test_paged_attention_kv_cache_prefill_and_decode(kv_cache_and_config):
kv_cache, rope_mode, support_sliding_window = kv_cache_and_config
if support_sliding_window and rope_mode == RopeMode.NORMAL:
diff --git
a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
index fe7fc0ff07..bd9b566c23 100644
--- a/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
+++ b/tests/python/relax/test_runtime_builtin_paged_attention_kv_cache_tir.py
@@ -18,7 +18,9 @@
import functools
import itertools
import json
+import math
+import numpy as np
import pytest
import torch
import tvm_ffi
@@ -91,6 +93,69 @@ fcompact_copy = None
_COMPILED_KERNEL_CACHE = {}
[email protected]
[email protected](not env.has_gpu(), reason="need gpu")
[email protected]("device_type", ["cuda", "metal"])
+def test_paged_prefill_layer_sliding_window_mask(device_type):
+ if not tvm.testing.device_enabled(device_type):
+ pytest.skip(f"{device_type} not enabled")
+
+ dev = tvm.device(device_type)
+ target = tvm.target.Target.from_device(dev)
+ head_dim = 64
+ num_kv_heads = 2
+ num_qo_heads = 4
+ page_size = 16
+ tir_func = _attention_prefill(
+ num_kv_heads,
+ num_qo_heads,
+ head_dim,
+ "float16",
+ True,
+ {},
+ target,
+ page_size=page_size,
+ sliding_window_size=3,
+ )
+ built = tvm.tirx.build(tir_func, target=target)
+
+ q = np.zeros((2, num_qo_heads, head_dim), dtype="float16")
+ q_indptr = np.array([0, 2], dtype="int32")
+ pages = np.zeros((1, 2, num_kv_heads, page_size, head_dim),
dtype="float16")
+ for position, value in enumerate([1, 3, 5]):
+ pages[0, 1, :, position, :] = value
+ page_indptr = np.array([0, 1], dtype="int32")
+ page_values = np.array([0], dtype="int32")
+ length_info = np.array([[3], [0], [0]], dtype="int32")
+ k_rope_pos_offset = np.array([0], dtype="int32")
+ q_rope_position = np.array([3, 4], dtype="int32")
+ output = np.zeros_like(q)
+ lse = np.zeros((2, num_qo_heads), dtype="float32")
+
+ args = [
+ tvm.runtime.tensor(array, device=dev)
+ for array in [
+ q,
+ q_indptr,
+ pages,
+ page_indptr,
+ page_values,
+ length_info,
+ k_rope_pos_offset,
+ q_rope_position,
+ output,
+ lse,
+ ]
+ ]
+
+ def run_and_check():
+ built.main(*args, 1, 0, 1.0, 10000.0, 1 / math.sqrt(head_dim))
+ expected = np.array([4, 5], dtype="float16")
+ tvm.testing.assert_allclose(args[8].numpy()[:, 0, 0], expected,
rtol=1e-3, atol=1e-3)
+
+ tvm.testing.run_with_gpu_lock(run_and_check)
+
+
def set_global_func(head_dim, dtype, target):
global fclear, fadd_sequence, fremove_sequence, ffork_sequence,
fenable_sliding_window_for_seq
global fpopn, fbegin_forward, fend_forward,
fcommit_accepted_token_tree_nodes
