From c18b75a20dd0494ac31fb0add3c7beb28ccf7819 Mon Sep 17 00:00:00 2001 From: zhangxin Date: Thu, 30 Apr 2026 15:54:36 +0800 Subject: [PATCH] =?UTF-8?q?issue/1148:=20PagedAttentionPrefill=20=E6=B7=BB?= =?UTF-8?q?=E5=8A=A0=20KV=20cache=20=E8=BF=9E=E7=BB=AD=E6=80=A7=20guard?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- python/infinicore/ops/paged_attention_prefill.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/infinicore/ops/paged_attention_prefill.py b/python/infinicore/ops/paged_attention_prefill.py index 848f74abf..259e865ae 100644 --- a/python/infinicore/ops/paged_attention_prefill.py +++ b/python/infinicore/ops/paged_attention_prefill.py @@ -2,6 +2,12 @@ from infinicore.tensor import Tensor +def _ensure_head_dim_contiguous(tensor: Tensor) -> Tensor: + if tensor.ndim > 0 and tensor.stride(tensor.ndim - 1) != 1: + return tensor.contiguous() + return tensor + + def paged_attention_prefill( q: Tensor, k_cache: Tensor, @@ -14,6 +20,8 @@ def paged_attention_prefill( *, out: Tensor | None = None, ): + k_cache = _ensure_head_dim_contiguous(k_cache) + v_cache = _ensure_head_dim_contiguous(v_cache) alibi_ptr = alibi_slopes._underlying if alibi_slopes is not None else None if out is None: