Skip to content

[DEV] nsa_paged_attention / nsa_compress_paged_cache 的 Metax 后端支持 #1393

Description

@Lfan-ke

目标版本

main

功能描述

nsa_paged_attentionnsa_compress_paged_cache 目前只有 cuda/ + nvidia/ 两层实现,缺少 metax/ 后端,因此这两个算子在 Metax GPU 上不可用。希望补齐它们的 Metax 后端,使 NSA(Native Sparse Attention)解码路径在 Metax 上可用。

实现要点

同族的 paged_attention 已经有完整的 Metax 后端,其 nvidiametax 实现之间只有机械差异(运行时头、devices/metax/*、命名空间、cudaStream_thcStream_tdevice::nvidia::Handledevice::metax::Handle),设备端 cuda/kernel.cuh 可原样复用。nsa_compress_paged_cache 与此完全同构。

nsa_paged_attention 多一个需要处理的点:它的 cuda/kernel.cuhblockReduceSum128 假定 warp 宽度为 32。

constexpr unsigned kFullMask = 0xffffffffu;
const int lane = threadIdx.x & 31;
const int warp = threadIdx.x >> 5;
for (int offset = 16; offset > 0; offset >>= 1) {
    value += __shfl_down_sync(kFullMask, value, offset);
}
__shared__ float warp_sums[4];

而 Metax 的 wavefront 是 64 lane(torch.cuda.get_device_properties(0).warp_size == 64,实测于 MetaX C500 / MACA 3.5.3),且 __shfl_down_sync 的掩码类型为 maca_uint64_t

__shfl_down_sync(const maca_uint64_t mask, const T var, const unsigned int delta, const int width = warpSize);

因此在 Metax 上,32 位的 0xffffffffu 掩码只会让 lane 0–31 参与归约,lane 32–63 的部分和被静默丢弃,块内求和结果错误;warp_sums[4]threadIdx.x >> 5 的分组也不再成立(128 线程在 Metax 上是 2 个 wavefront 而非 4 个)。这里打算按仓库已有的 ENABLE_METAX_API 惯例(见 var/cuda/kernel.cuhtopk/cuda/kernel.cuh)把 warp 宽度、掩码与 warp 数抽成常量,使该归约对两种 wavefront 宽度都成立,不改变 NVIDIA 路径的行为。

会在 MetaX C500 上跑 test/infiniop/nsa_paged_attention.py --metaxtest/infiniop/nsa_compress_paged_cache.py --metax 并附通过截图。

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