Skip to content

[BUG] PagedAttentionPrefill KV cache head_dim 维度 stride 不连续导致底层算子报错 #1148

@JoeZhang-0x000

Description

@JoeZhang-0x000

问题描述

PagedAttentionPrefill 底层算子要求 KV cache 最后一维(head_dim)的 stride 必须为 1,但当前实现未对传入的 k_cache / v_cache 做连续性检查,导致非连续张量直接传入底层算子时触发 Bad Tensor Strides 错误。该问题是普遍性的,不限于特定后端。

复现步骤

python test/infinicore/ops/paged_attention_prefill.py --metax

结果:failed 56/60

错误信息

Bad Tensor Strides
k_cache_desc->stride(3) == 1

传入张量 stride 类似 (4096, 1, 512, 4),最后一维 stride 为 4,不满足 stride(3) == 1 的要求。

建议修复

python/infinicore/ops/paged_attention_prefill.py 的 Python wrapper 中添加连续性 guard:当 k_cache / v_cache 的最后一维不是连续时,自动调用 .contiguous() 转换。

标准 vLLM KV cache view 的 head-dim stride 已经是 1,所以正常 vLLM 路径不会触发额外 copy。

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions