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 4d3f3dacf2 [Runtime] Clear strides when the paged KV cache copies one
row of aux data (#20491)
4d3f3dacf2 is described below
commit 4d3f3dacf29f6d7878bd7d0b2bab0d584a7f869b
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Tue Sep 29 18:10:37 2026 -0700
[Runtime] Clear strides when the paged KV cache copies one row of aux data
(#20491)
`PlainPagedKVCacheAuxDataManager` fills the length info and the compact
copy positions one row at a time into a `{rows, n}` view.
`CopyVecDataToArray` narrows the destination to one dimension for each
row but keeps the strides of the two-dimensional view, so
`CopyDataFromTo` rejects it as non-contiguous once `n` is above one.
With sliding window support enabled that is any forward pass over more
than one sequence, on every device that uses the plain manager (Metal,
Vulkan, OpenCL, WebGPU). MLC LLM hits it when serving concurrent Gemma 4
requests on Metal.
1. Clear the strides together with the shape in `CopyVecDataToArray`, as
the cached manager already does
2. Add a test that runs two sequences through a cache with sliding
window support on CUDA and Metal
---
src/runtime/vm/attn_utils.h | 2 ++
...runtime_builtin_paged_attention_kv_cache_tir.py | 35 ++++++++++++++++++++++
2 files changed, 37 insertions(+)
diff --git a/src/runtime/vm/attn_utils.h b/src/runtime/vm/attn_utils.h
index dba893d4d4..3bc8dc5174 100644
--- a/src/runtime/vm/attn_utils.h
+++ b/src/runtime/vm/attn_utils.h
@@ -768,6 +768,8 @@ class PlainPagedKVCacheAuxDataManager : public
PagedKVCacheAuxDataManager {
TVM_FFI_ICHECK_EQ(shape.value().size(), 1);
copy_dst.ndim = 1;
copy_dst.shape = const_cast<int64_t*>(shape.value()->data);
+ // The strides still describe the multi-row view this copies one row of.
+ copy_dst.strides = nullptr;
}
copy_dst.byte_offset = dst_elem_offset * sizeof(int32_t);
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 bd9b566c23..7f878f118f 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
@@ -156,6 +156,41 @@ def
test_paged_prefill_layer_sliding_window_mask(device_type):
tvm.testing.run_with_gpu_lock(run_and_check)
[email protected]
[email protected](not env.has_gpu(), reason="need gpu")
[email protected]("device_type", ["cuda", "metal"])
+def test_paged_attention_kv_cache_sliding_window_two_sequences(device_type):
+ if not tvm.testing.device_enabled(device_type):
+ pytest.skip(f"{device_type} not enabled")
+
+ head_dim = 64
+ dtype = "float16"
+
+ def run_and_check():
+ global device
+ device = tvm.device(device_type)
+ set_global_func(head_dim, dtype, tvm.target.Target.from_device(device))
+ try:
+ cache = create_kv_cache(head_dim, dtype, RopeMode.NONE, True)
+ fadd_sequence(cache, 0)
+ fadd_sequence(cache, 1)
+ fbegin_forward(cache, Shape([0, 1]), Shape([1, 1]))
+ qkv = np.random.rand(2, num_qo_heads + 2 * num_kv_heads,
head_dim).astype(dtype)
+ output = tvm.runtime.empty((2, num_qo_heads, head_dim), dtype,
device=device)
+ fattention_with_fuse_qkv(
+ cache, 0, head_dim**-0.5, tvm.runtime.tensor(qkv,
device=device), output
+ )
+ fend_forward(cache)
+ finally:
+ device = None
+ # Each sequence holds one token, so every query head returns the value
of its KV head.
+ values = qkv[:, num_qo_heads + num_kv_heads :, :]
+ expected = np.repeat(values, num_qo_heads // num_kv_heads, axis=1)
+ tvm.testing.assert_allclose(output.numpy(), 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
